Add block_range_read_stream_cb(), to deduplicate code.

nmisch · nmisch · commit c582b75851c2 · 2024-09-03T10:46:20.000-07:00
This replaces two functions for iterating over all blocks in a range. A pending patch will use this instead of adding a third. Nazir Bilal Yavuz Discussion: https://postgr.es/m/20240820184742.f2.nmisch@google.com
diff --git a/contrib/pg_prewarm/pg_prewarm.c b/contrib/pg_prewarm/pg_prewarm.c
@@ -39,25 +39,6 @@ typedef enum
 
 static PGIOAlignedBlock blockbuffer;
 
-struct pg_prewarm_read_stream_private
-{
-	BlockNumber blocknum;
-	int64		last_block;
-};
-
-static BlockNumber
-pg_prewarm_read_stream_next_block(ReadStream *stream,
-								  void *callback_private_data,
-								  void *per_buffer_data)
-{
-	struct pg_prewarm_read_stream_private *p = callback_private_data;
-
-	if (p->blocknum <= p->last_block)
-		return p->blocknum++;
-
-	return InvalidBlockNumber;
-}
-
 /*
  * pg_prewarm(regclass, mode text, fork text,
  *			  first_block int8, last_block int8)
@@ -203,22 +184,22 @@ pg_prewarm(PG_FUNCTION_ARGS)
 	}
 	else if (ptype == PREWARM_BUFFER)
 	{
-		struct pg_prewarm_read_stream_private p;
+		BlockRangeReadStreamPrivate p;
 		ReadStream *stream;
 
 		/*
 		 * In buffer mode, we actually pull the data into shared_buffers.
 		 */
 
 		/* Set up the private state for our streaming buffer read callback. */
-		p.blocknum = first_block;
-		p.last_block = last_block;
+		p.current_blocknum = first_block;
+		p.last_exclusive = last_block + 1;
 
 		stream = read_stream_begin_relation(READ_STREAM_FULL,
 											NULL,
 											rel,
 											forkNumber,
-											pg_prewarm_read_stream_next_block,
+											block_range_read_stream_cb,
 											&p,
 											0);
 
diff --git a/src/backend/storage/aio/read_stream.c b/src/backend/storage/aio/read_stream.c
@@ -163,6 +163,23 @@ get_per_buffer_data(ReadStream *stream, int16 buffer_index)
 		stream->per_buffer_data_size * buffer_index;
 }
 
+/*
+ * General-use ReadStreamBlockNumberCB for block range scans.  Loops over the
+ * blocks [current_blocknum, last_exclusive).
+ */
+BlockNumber
+block_range_read_stream_cb(ReadStream *stream,
+						   void *callback_private_data,
+						   void *per_buffer_data)
+{
+	BlockRangeReadStreamPrivate *p = callback_private_data;
+
+	if (p->current_blocknum < p->last_exclusive)
+		return p->current_blocknum++;
+
+	return InvalidBlockNumber;
+}
+
 /*
  * Ask the callback which block it would like us to read next, with a one block
  * buffer in front to allow read_stream_unget_block() to work.
diff --git a/src/backend/storage/buffer/bufmgr.c b/src/backend/storage/buffer/bufmgr.c
@@ -136,33 +136,6 @@ typedef struct SMgrSortArray
 	SMgrRelation srel;
 } SMgrSortArray;
 
-/*
- * Helper struct for read stream object used in
- * RelationCopyStorageUsingBuffer() function.
- */
-struct copy_storage_using_buffer_read_stream_private
-{
-	BlockNumber blocknum;
-	BlockNumber nblocks;
-};
-
-/*
- * Callback function to get next block for read stream object used in
- * RelationCopyStorageUsingBuffer() function.
- */
-static BlockNumber
-copy_storage_using_buffer_read_stream_next_block(ReadStream *stream,
-												 void *callback_private_data,
-												 void *per_buffer_data)
-{
-	struct copy_storage_using_buffer_read_stream_private *p = callback_private_data;
-
-	if (p->blocknum < p->nblocks)
-		return p->blocknum++;
-
-	return InvalidBlockNumber;
-}
-
 /* GUC variables */
 bool		zero_damaged_pages = false;
 int			bgwriter_lru_maxpages = 100;
@@ -4710,7 +4683,7 @@ RelationCopyStorageUsingBuffer(RelFileLocator srclocator,
 	PGIOAlignedBlock buf;
 	BufferAccessStrategy bstrategy_src;
 	BufferAccessStrategy bstrategy_dst;
-	struct copy_storage_using_buffer_read_stream_private p;
+	BlockRangeReadStreamPrivate p;
 	ReadStream *src_stream;
 	SMgrRelation src_smgr;
 
@@ -4742,15 +4715,15 @@ RelationCopyStorageUsingBuffer(RelFileLocator srclocator,
 	bstrategy_dst = GetAccessStrategy(BAS_BULKWRITE);
 
 	/* Initialize streaming read */
-	p.blocknum = 0;
-	p.nblocks = nblocks;
+	p.current_blocknum = 0;
+	p.last_exclusive = nblocks;
 	src_smgr = smgropen(srclocator, INVALID_PROC_NUMBER);
 	src_stream = read_stream_begin_smgr_relation(READ_STREAM_FULL,
 												 bstrategy_src,
 												 src_smgr,
 												 permanent ? RELPERSISTENCE_PERMANENT : RELPERSISTENCE_UNLOGGED,
 												 forkNum,
-												 copy_storage_using_buffer_read_stream_next_block,
+												 block_range_read_stream_cb,
 												 &p,
 												 0);
 
diff --git a/src/include/storage/read_stream.h b/src/include/storage/read_stream.h
@@ -45,11 +45,21 @@
 struct ReadStream;
 typedef struct ReadStream ReadStream;
 
+/* for block_range_read_stream_cb */
+typedef struct BlockRangeReadStreamPrivate
+{
+	BlockNumber current_blocknum;
+	BlockNumber last_exclusive;
+} BlockRangeReadStreamPrivate;
+
 /* Callback that returns the next block number to read. */
 typedef BlockNumber (*ReadStreamBlockNumberCB) (ReadStream *stream,
 												void *callback_private_data,
 												void *per_buffer_data);
 
+extern BlockNumber block_range_read_stream_cb(ReadStream *stream,
+											  void *callback_private_data,
+											  void *per_buffer_data);
 extern ReadStream *read_stream_begin_relation(int flags,
 											  BufferAccessStrategy strategy,
 											  Relation rel,
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
@@ -275,6 +275,7 @@ BlockId
 BlockIdData
 BlockInfoRecord
 BlockNumber
+BlockRangeReadStreamPrivate
 BlockRefTable
 BlockRefTableBuffer
 BlockRefTableChunk
