v1-0008-WIP-Use-streaming-reads-in-heapam-scans.patch

From 091bc69a8b02e755d55b27d58fb1c17e9ab9a6dc Mon Sep 17 00:00:00 2001
From: Andres Freund <andres@anarazel.de>
Date: Wed, 12 Apr 2023 18:16:06 -0700
Subject: [PATCH v1 08/14] WIP: Use streaming reads in heapam scans.

XXX Cherry-picked from https://github.com/anarazel/postgres/tree/aio and
lightly modified by TM, for demonstration purposes.

Author: Andres Freund <andres@anarazel.de>
---
 src/backend/access/heap/heapam.c         | 220 +++++++++++++++++++++--
 src/backend/access/heap/heapam_handler.c |   2 +-
 src/include/access/heapam.h              |   5 +-
 3 files changed, 207 insertions(+), 20 deletions(-)

diff --git a/src/backend/access/heap/heapam.c b/src/backend/access/heap/heapam.c
index 6a66214a58..01cca53586 100644
--- a/src/backend/access/heap/heapam.c
+++ b/src/backend/access/heap/heapam.c
@@ -65,6 +65,7 @@
 #include "storage/smgr.h"
 #include "storage/spin.h"
 #include "storage/standby.h"
+#include "storage/streaming_read.h"
 #include "utils/datum.h"
 #include "utils/inval.h"
 #include "utils/lsyscache.h"
@@ -225,6 +226,99 @@ static const int MultiXactStatusLock[MaxMultiXactStatus + 1] =
  * ----------------------------------------------------------------
  */
 
+static bool
+heap_pgsr_next_single(PgStreamingRead *pgsr,
+					  uintptr_t pgsr_private, void *per_io_data,
+					  BufferManagerRelation *bmr, ForkNumber *fork,
+					  BlockNumber *block, ReadBufferMode *mode)
+{
+	HeapScanDesc scan = (HeapScanDesc) pgsr_private;
+	BlockNumber blockno;
+
+	Assert(!scan->rs_base.rs_parallel);
+	Assert(scan->rs_nblocks > 0);
+
+	if (scan->rs_prefetch_block == InvalidBlockNumber)
+	{
+		scan->rs_prefetch_block = blockno = scan->rs_startblock;
+	}
+	else
+	{
+		blockno = ++scan->rs_prefetch_block;
+
+		/* wrap back to the start of the heap */
+		if (blockno >= scan->rs_nblocks)
+			scan->rs_prefetch_block = blockno = 0;
+
+		/* we're done if we're back at where we started */
+		if (blockno == scan->rs_startblock)
+			return false;
+
+		/* check if the limit imposed by heap_setscanlimits() is met */
+		if (scan->rs_numblocks != InvalidBlockNumber)
+		{
+			if (--scan->rs_numblocks == 0)
+				return false;
+		}
+	}
+
+	*bmr = BMR_REL(scan->rs_base.rs_rd);
+	*fork = MAIN_FORKNUM;
+	*block = blockno;
+	*mode = RBM_NORMAL;
+
+	return true;
+}
+
+static bool
+heap_pgsr_next_parallel(PgStreamingRead *pgsr,
+						uintptr_t pgsr_private, void *per_io_data,
+						BufferManagerRelation *bmr, ForkNumber *fork,
+						BlockNumber *block, ReadBufferMode *mode)
+{
+	HeapScanDesc scan = (HeapScanDesc) pgsr_private;
+	ParallelBlockTableScanDesc pbscan =
+		(ParallelBlockTableScanDesc) scan->rs_base.rs_parallel;
+	ParallelBlockTableScanWorker pbscanwork =
+		scan->rs_parallelworkerdata;
+	BlockNumber blockno;
+
+	Assert(scan->rs_base.rs_parallel);
+	Assert(scan->rs_nblocks > 0);
+
+	/* Note that other processes might have already finished the scan */
+	blockno = table_block_parallelscan_nextpage(scan->rs_base.rs_rd,
+												pbscanwork, pbscan);
+	if (blockno == InvalidBlockNumber)
+		return false;
+
+	*bmr = BMR_REL(scan->rs_base.rs_rd);
+	*fork = MAIN_FORKNUM;
+	*block = blockno;
+	*mode = RBM_NORMAL;
+
+	return true;
+}
+
+static PgStreamingRead *
+heap_pgsr_single_alloc(HeapScanDesc scan)
+{
+	return pg_streaming_read_buffer_alloc(PG_STREAMING_READ_DEFAULT_MAX_IOS,
+										  0, (uintptr_t) scan,
+										  scan->rs_strategy,
+										  heap_pgsr_next_single);
+}
+
+static PgStreamingRead *
+heap_pgsr_parallel_alloc(HeapScanDesc scan)
+{
+	return pg_streaming_read_buffer_alloc(PG_STREAMING_READ_DEFAULT_MAX_IOS,
+										  0, (uintptr_t) scan,
+										  scan->rs_strategy,
+										  heap_pgsr_next_parallel);
+}
+
+
 /* ----------------
  *		initscan - scan code common to heap_beginscan and heap_rescan
  * ----------------
@@ -342,6 +436,26 @@ initscan(HeapScanDesc scan, ScanKey key, bool keep_startblock)
 	 */
 	if (scan->rs_base.rs_flags & SO_TYPE_SEQSCAN)
 		pgstat_count_heap_scan(scan->rs_base.rs_rd);
+
+	scan->rs_prefetch_block = InvalidBlockNumber;
+	if (scan->pgsr)
+	{
+		pg_streaming_read_free(scan->pgsr);
+		scan->pgsr = NULL;
+	}
+
+	/*
+	 * FIXME: This probably should be done in the !rs_inited blocks instead.
+	 */
+	scan->pgsr = NULL;
+	if (!RelationUsesLocalBuffers(scan->rs_base.rs_rd) &&
+		(scan->rs_base.rs_flags & SO_TYPE_SEQSCAN))
+	{
+		if (scan->rs_base.rs_parallel)
+			scan->pgsr = heap_pgsr_parallel_alloc(scan);
+		else
+			scan->pgsr = heap_pgsr_single_alloc(scan);
+	}
 }
 
 /*
@@ -374,7 +488,7 @@ heap_setscanlimits(TableScanDesc sscan, BlockNumber startBlk, BlockNumber numBlk
  * which tuples on the page are visible.
  */
 void
-heapgetpage(TableScanDesc sscan, BlockNumber block)
+heapgetpage(TableScanDesc sscan, BlockNumber block, Buffer pgsr_buffer)
 {
 	HeapScanDesc scan = (HeapScanDesc) sscan;
 	Buffer		buffer;
@@ -401,9 +515,20 @@ heapgetpage(TableScanDesc sscan, BlockNumber block)
 	 */
 	CHECK_FOR_INTERRUPTS();
 
-	/* read page using selected strategy */
-	scan->rs_cbuf = ReadBufferExtended(scan->rs_base.rs_rd, MAIN_FORKNUM, block,
-									   RBM_NORMAL, scan->rs_strategy);
+	if (BufferIsValid(pgsr_buffer))
+	{
+		Assert(scan->pgsr);
+		Assert(BufferGetBlockNumber(pgsr_buffer) == block);
+		scan->rs_cbuf = pgsr_buffer;
+	}
+	else
+	{
+		Assert(!scan->pgsr);
+
+		/* read page using selected strategy */
+		scan->rs_cbuf = ReadBufferExtended(scan->rs_base.rs_rd, MAIN_FORKNUM, block,
+										   RBM_NORMAL, scan->rs_strategy);
+	}
 	scan->rs_cblock = block;
 
 	if (!(scan->rs_base.rs_flags & SO_ALLOW_PAGEMODE))
@@ -491,7 +616,7 @@ heapgetpage(TableScanDesc sscan, BlockNumber block)
  * of the pages before we can get a chance to get our first page.
  */
 static BlockNumber
-heapgettup_initial_block(HeapScanDesc scan, ScanDirection dir)
+heapgettup_initial_block(HeapScanDesc scan, ScanDirection dir, Buffer *pgsr_buf)
 {
 	Assert(!scan->rs_inited);
 
@@ -501,16 +626,25 @@ heapgettup_initial_block(HeapScanDesc scan, ScanDirection dir)
 
 	if (ScanDirectionIsForward(dir))
 	{
+		if (scan->rs_base.rs_parallel != NULL)
+			table_block_parallelscan_startblock_init(scan->rs_base.rs_rd,
+													 scan->rs_parallelworkerdata,
+													 (ParallelBlockTableScanDesc) scan->rs_base.rs_parallel);
+
+		/* FIXME: Integrate more neatly */
+		if (scan->pgsr)
+		{
+			*pgsr_buf = pg_streaming_read_buffer_get_next(scan->pgsr, NULL);
+			if (*pgsr_buf == InvalidBuffer)
+				return InvalidBlockNumber;
+			return BufferGetBlockNumber(*pgsr_buf);
+		}
+
 		/* serial scan */
 		if (scan->rs_base.rs_parallel == NULL)
 			return scan->rs_startblock;
 		else
 		{
-			/* parallel scan */
-			table_block_parallelscan_startblock_init(scan->rs_base.rs_rd,
-													 scan->rs_parallelworkerdata,
-													 (ParallelBlockTableScanDesc) scan->rs_base.rs_parallel);
-
 			/* may return InvalidBlockNumber if there are no more blocks */
 			return table_block_parallelscan_nextpage(scan->rs_base.rs_rd,
 													 scan->rs_parallelworkerdata,
@@ -530,6 +664,12 @@ heapgettup_initial_block(HeapScanDesc scan, ScanDirection dir)
 		 */
 		scan->rs_base.rs_flags &= ~SO_ALLOW_SYNC;
 
+		if (scan->pgsr)
+		{
+			pg_streaming_read_free(scan->pgsr);
+			scan->pgsr = NULL;
+		}
+
 		/*
 		 * Start from last page of the scan.  Ensure we take into account
 		 * rs_numblocks if it's been adjusted by heap_setscanlimits().
@@ -635,11 +775,38 @@ heapgettup_continue_page(HeapScanDesc scan, ScanDirection dir, int *linesleft,
  * heap_setscanlimits().
  */
 static inline BlockNumber
-heapgettup_advance_block(HeapScanDesc scan, BlockNumber block, ScanDirection dir)
+heapgettup_advance_block(HeapScanDesc scan, BlockNumber block, ScanDirection dir,
+						 Buffer *pgsr_buf)
 {
 	if (ScanDirectionIsForward(dir))
 	{
-		if (scan->rs_base.rs_parallel == NULL)
+		if (scan->pgsr)
+		{
+#ifdef USE_ASSERT_CHECKING
+			block++;
+
+			/* wrap back to the start of the heap */
+			if (block >= scan->rs_nblocks)
+				block = 0;
+
+			/* we're done if we're back at where we started */
+			if (block == scan->rs_startblock)
+				block = InvalidBlockNumber;
+			else if (scan->rs_numblocks != InvalidBlockNumber)
+			{
+				if (--scan->rs_numblocks == 0)
+					block = InvalidBlockNumber;
+			}
+#endif
+			*pgsr_buf = pg_streaming_read_buffer_get_next(scan->pgsr, NULL);
+			if (*pgsr_buf == InvalidBuffer)
+				return InvalidBlockNumber;
+
+			Assert(scan->rs_base.rs_parallel ||
+				   block == BufferGetBlockNumber(*pgsr_buf));
+			return BufferGetBlockNumber(*pgsr_buf);
+		}
+		else if (scan->rs_base.rs_parallel == NULL)
 		{
 			block++;
 
@@ -684,6 +851,12 @@ heapgettup_advance_block(HeapScanDesc scan, BlockNumber block, ScanDirection dir
 	}
 	else
 	{
+		if (scan->pgsr)
+		{
+			pg_streaming_read_free(scan->pgsr);
+			scan->pgsr = NULL;
+		}
+
 		/* we're done if the last block is the start position */
 		if (block == scan->rs_startblock)
 			return InvalidBlockNumber;
@@ -733,13 +906,14 @@ heapgettup(HeapScanDesc scan,
 {
 	HeapTuple	tuple = &(scan->rs_ctup);
 	BlockNumber block;
+	Buffer		pgsr_buf = InvalidBuffer;
 	Page		page;
 	OffsetNumber lineoff;
 	int			linesleft;
 
 	if (unlikely(!scan->rs_inited))
 	{
-		block = heapgettup_initial_block(scan, dir);
+		block = heapgettup_initial_block(scan, dir, &pgsr_buf);
 		/* ensure rs_cbuf is invalid when we get InvalidBlockNumber */
 		Assert(block != InvalidBlockNumber || !BufferIsValid(scan->rs_cbuf));
 		scan->rs_inited = true;
@@ -760,7 +934,7 @@ heapgettup(HeapScanDesc scan,
 	 */
 	while (block != InvalidBlockNumber)
 	{
-		heapgetpage((TableScanDesc) scan, block);
+		heapgetpage((TableScanDesc) scan, block, pgsr_buf);
 		LockBuffer(scan->rs_cbuf, BUFFER_LOCK_SHARE);
 		page = heapgettup_start_page(scan, dir, &linesleft, &lineoff);
 continue_page:
@@ -812,9 +986,10 @@ continue_page:
 		 * it's time to move to the next.
 		 */
 		LockBuffer(scan->rs_cbuf, BUFFER_LOCK_UNLOCK);
+		pgsr_buf = InvalidBuffer;
 
 		/* get the BlockNumber to scan next */
-		block = heapgettup_advance_block(scan, block, dir);
+		block = heapgettup_advance_block(scan, block, dir, &pgsr_buf);
 	}
 
 	/* end of scan */
@@ -848,13 +1023,14 @@ heapgettup_pagemode(HeapScanDesc scan,
 {
 	HeapTuple	tuple = &(scan->rs_ctup);
 	BlockNumber block;
+	Buffer		pgsr_buf = InvalidBuffer;
 	Page		page;
 	int			lineindex;
 	int			linesleft;
 
 	if (unlikely(!scan->rs_inited))
 	{
-		block = heapgettup_initial_block(scan, dir);
+		block = heapgettup_initial_block(scan, dir, &pgsr_buf);
 		/* ensure rs_cbuf is invalid when we get InvalidBlockNumber */
 		Assert(block != InvalidBlockNumber || !BufferIsValid(scan->rs_cbuf));
 		scan->rs_inited = true;
@@ -882,7 +1058,7 @@ heapgettup_pagemode(HeapScanDesc scan,
 	 */
 	while (block != InvalidBlockNumber)
 	{
-		heapgetpage((TableScanDesc) scan, block);
+		heapgetpage((TableScanDesc) scan, block, pgsr_buf);
 		page = BufferGetPage(scan->rs_cbuf);
 		TestForOldSnapshot(scan->rs_base.rs_snapshot, scan->rs_base.rs_rd, page);
 		linesleft = scan->rs_ntuples;
@@ -915,7 +1091,7 @@ continue_page:
 		}
 
 		/* get the BlockNumber to scan next */
-		block = heapgettup_advance_block(scan, block, dir);
+		block = heapgettup_advance_block(scan, block, dir, &pgsr_buf);
 	}
 
 	/* end of scan */
@@ -963,6 +1139,8 @@ heap_beginscan(Relation relation, Snapshot snapshot,
 	scan->rs_base.rs_parallel = parallel_scan;
 	scan->rs_strategy = NULL;	/* set in initscan */
 
+	scan->pgsr = NULL;
+
 	/*
 	 * Disable page-at-a-time mode if it's not a MVCC-safe snapshot.
 	 */
@@ -1069,6 +1247,12 @@ heap_endscan(TableScanDesc sscan)
 	if (BufferIsValid(scan->rs_cbuf))
 		ReleaseBuffer(scan->rs_cbuf);
 
+	if (scan->pgsr)
+	{
+		pg_streaming_read_free(scan->pgsr);
+		scan->pgsr = NULL;
+	}
+
 	/*
 	 * decrement relation reference count and free scan descriptor storage
 	 */
diff --git a/src/backend/access/heap/heapam_handler.c b/src/backend/access/heap/heapam_handler.c
index 5a17112c91..26378cbb4b 100644
--- a/src/backend/access/heap/heapam_handler.c
+++ b/src/backend/access/heap/heapam_handler.c
@@ -2339,7 +2339,7 @@ heapam_scan_sample_next_block(TableScanDesc scan, SampleScanState *scanstate)
 		return false;
 	}
 
-	heapgetpage(scan, blockno);
+	heapgetpage(scan, blockno, InvalidBuffer);
 	hscan->rs_inited = true;
 
 	return true;
diff --git a/src/include/access/heapam.h b/src/include/access/heapam.h
index faf5026519..852d47b1f8 100644
--- a/src/include/access/heapam.h
+++ b/src/include/access/heapam.h
@@ -59,6 +59,7 @@ typedef struct HeapScanDescData
 	bool		rs_inited;		/* false = scan not init'd yet */
 	OffsetNumber rs_coffset;	/* current offset # in non-page-at-a-time mode */
 	BlockNumber rs_cblock;		/* current block # in scan, if any */
+	BlockNumber rs_prefetch_block;	/* block being prefetched */
 	Buffer		rs_cbuf;		/* current buffer in scan, if any */
 	/* NB: if rs_cbuf is not InvalidBuffer, we hold a pin on that buffer */
 
@@ -72,6 +73,8 @@ typedef struct HeapScanDescData
 	 */
 	ParallelBlockTableScanWorkerData *rs_parallelworkerdata;
 
+	struct PgStreamingRead *pgsr;
+
 	/* these fields only used in page-at-a-time mode and for bitmap scans */
 	int			rs_cindex;		/* current tuple's index in vistuples */
 	int			rs_ntuples;		/* number of visible tuples on page */
@@ -212,7 +215,7 @@ extern TableScanDesc heap_beginscan(Relation relation, Snapshot snapshot,
 									uint32 flags);
 extern void heap_setscanlimits(TableScanDesc sscan, BlockNumber startBlk,
 							   BlockNumber numBlks);
-extern void heapgetpage(TableScanDesc sscan, BlockNumber block);
+extern void heapgetpage(TableScanDesc sscan, BlockNumber block, Buffer buffer);
 extern void heap_rescan(TableScanDesc sscan, ScanKey key, bool set_params,
 						bool allow_strat, bool allow_sync, bool allow_pagemode);
 extern void heap_endscan(TableScanDesc sscan);
-- 
2.39.2