[CF 5617] v2 - Beautify read stream "per buffer data" APIs

This branch was automatically generated by a robot using patches from an
email thread registered at:

https://commitfest.postgresql.org/patch/5617

The branch will be overwritten each time a new patch version is posted to
the thread, and also periodically to check for bitrot caused by changes
on the master branch.

Patch(es): https://www.postgresql.org/message-id/CA+hUKGLa7ba7USyT+JR7uRiawWeCVJ96wyRsoEXk7r2gngPv=A@mail.gmail.com
Author(s): Thomas Munro

Author: Commitfest Bot

contrib/pg_prewarm/pg_prewarm.c

@@ -216,11 +216,11 @@ pg_prewarm(PG_FUNCTION_ARGS)
 			Buffer		buf;
 
 			CHECK_FOR_INTERRUPTS();
-			buf = read_stream_next_buffer(stream, NULL);
+			buf = read_stream_get_buffer(stream);
 			ReleaseBuffer(buf);
 			++blocks_done;
 		}
-		Assert(read_stream_next_buffer(stream, NULL) == InvalidBuffer);
+		Assert(read_stream_get_buffer(stream) == InvalidBuffer);
 		read_stream_end(stream);
 	}
 

contrib/pg_visibility/pg_visibility.c

@@ -565,7 +565,7 @@ collect_visibility_data(Oid relid, bool include_pd)
 			Buffer		buffer;
 			Page		page;
 
-			buffer = read_stream_next_buffer(stream, NULL);
+			buffer = read_stream_get_buffer(stream);
 			LockBuffer(buffer, BUFFER_LOCK_SHARE);
 
 			page = BufferGetPage(buffer);
@@ -578,7 +578,7 @@ collect_visibility_data(Oid relid, bool include_pd)
 
 	if (include_pd)
 	{
-		Assert(read_stream_next_buffer(stream, NULL) == InvalidBuffer);
+		Assert(read_stream_get_buffer(stream) == InvalidBuffer);
 		read_stream_end(stream);
 	}
 
@@ -761,7 +761,7 @@ collect_corrupt_items(Oid relid, bool all_visible, bool all_frozen)
 										0);
 
 	/* Loop over every block in the relation. */
-	while ((buffer = read_stream_next_buffer(stream, NULL)) != InvalidBuffer)
+	while ((buffer = read_stream_get_buffer(stream)) != InvalidBuffer)
 	{
 		bool		check_frozen = all_frozen;
 		bool		check_visible = all_visible;

src/backend/access/heap/heapam.c

@@ -651,7 +651,7 @@ heap_fetch_next_buffer(HeapScanDesc scan, ScanDirection dir)
 
 	scan->rs_dir = dir;
 
-	scan->rs_cbuf = read_stream_next_buffer(scan->rs_read_stream, NULL);
+	scan->rs_cbuf = read_stream_get_buffer(scan->rs_read_stream);
 	if (BufferIsValid(scan->rs_cbuf))
 		scan->rs_cblock = BufferGetBlockNumber(scan->rs_cbuf);
 }

src/backend/access/heap/heapam_handler.c

@@ -1014,7 +1014,7 @@ heapam_scan_analyze_next_block(TableScanDesc scan, ReadStream *stream)
 	 * re-acquire sharelock for each tuple, but since we aren't doing much
 	 * work per tuple, the extra lock traffic is probably better avoided.
 	 */
-	hscan->rs_cbuf = read_stream_next_buffer(stream, NULL);
+	hscan->rs_cbuf = read_stream_get_buffer(stream);
 	if (!BufferIsValid(hscan->rs_cbuf))
 		return false;
 

src/backend/access/heap/vacuumlazy.c

@@ -1245,7 +1245,6 @@ lazy_scan_heap(LVRelState *vacrel)
 		Page		page;
 		uint8		blk_info = 0;
 		bool		has_lpdead_items;
-		void	   *per_buffer_data = NULL;
 		bool		vm_page_frozen = false;
 		bool		got_cleanup_lock = false;
 
@@ -1305,13 +1304,12 @@ lazy_scan_heap(LVRelState *vacrel)
 										 PROGRESS_VACUUM_PHASE_SCAN_HEAP);
 		}
 
-		buf = read_stream_next_buffer(stream, &per_buffer_data);
+		buf = read_stream_get_buffer_and_value(stream, &blk_info);
 
 		/* The relation is exhausted. */
 		if (!BufferIsValid(buf))
 			break;
 
-		blk_info = *((uint8 *) per_buffer_data);
 		CheckBufferIsPinnedOnce(buf);
 		page = BufferGetPage(buf);
 		blkno = BufferGetBlockNumber(buf);
@@ -1632,7 +1630,7 @@ heap_vac_scan_next_block(ReadStream *stream,
 		 */
 		vacrel->current_block = next_block;
 		blk_info |= VAC_BLK_ALL_VISIBLE_ACCORDING_TO_VM;
-		*((uint8 *) per_buffer_data) = blk_info;
+		read_stream_put_value(stream, per_buffer_data, blk_info);
 		return vacrel->current_block;
 	}
 	else
@@ -1648,7 +1646,7 @@ heap_vac_scan_next_block(ReadStream *stream,
 			blk_info |= VAC_BLK_ALL_VISIBLE_ACCORDING_TO_VM;
 		if (vacrel->next_unskippable_eager_scanned)
 			blk_info |= VAC_BLK_WAS_EAGER_SCANNED;
-		*((uint8 *) per_buffer_data) = blk_info;
+		read_stream_put_value(stream, per_buffer_data, blk_info);
 		return vacrel->current_block;
 	}
 }
@@ -2693,7 +2691,7 @@ vacuum_reap_lp_read_stream_next(ReadStream *stream,
 	 * Save the TidStoreIterResult for later, so we can extract the offsets.
 	 * It is safe to copy the result, according to TidStoreIterateNext().
 	 */
-	memcpy(per_buffer_data, iter_result, sizeof(*iter_result));
+	read_stream_put_value(stream, per_buffer_data, *iter_result);
 
 	return iter_result->blkno;
 }
@@ -2768,7 +2766,7 @@ lazy_vacuum_heap_rel(LVRelState *vacrel)
 
 		vacuum_delay_point(false);
 
-		buf = read_stream_next_buffer(stream, (void **) &iter_result);
+		buf = read_stream_get_buffer_and_pointer(stream, &iter_result);
 
 		/* The relation is exhausted */
 		if (!BufferIsValid(buf))

src/backend/storage/aio/read_stream.c

@@ -762,6 +762,9 @@ read_stream_begin_smgr_relation(int flags,
  * valid until the next call to read_stream_next_buffer().  When the stream
  * runs out of data, InvalidBuffer is returned.  The caller may decide to end
  * the stream early at any time by calling read_stream_end().
+ *
+ * See read_stream.h for read_stream_get_buffer() and variants that provide
+ * some degree of type safety for the per_buffer_data argument.
  */
 Buffer
 read_stream_next_buffer(ReadStream *stream, void **per_buffer_data)
@@ -997,6 +1000,15 @@ read_stream_next_block(ReadStream *stream, BufferAccessStrategy *strategy)
 	return read_stream_get_block(stream, NULL);
 }
 
+/*
+ * Return the configured per-buffer data size, for use in assertions.
+ */
+size_t
+read_stream_per_buffer_data_size(ReadStream *stream)
+{
+	return stream->per_buffer_data_size;
+}
+
 /*
  * Reset a read stream by releasing any queued up buffers, allowing the stream
  * to be used again for different blocks.  This can be used to clear an

src/backend/storage/buffer/bufmgr.c

@@ -5117,7 +5117,7 @@ RelationCopyStorageUsingBuffer(RelFileLocator srclocator,
 		CHECK_FOR_INTERRUPTS();
 
 		/* Read block from source relation. */
-		srcBuf = read_stream_next_buffer(src_stream, NULL);
+		srcBuf = read_stream_get_buffer(src_stream);
 		LockBuffer(srcBuf, BUFFER_LOCK_SHARE);
 		srcPage = BufferGetPage(srcBuf);
 
@@ -5142,7 +5142,7 @@ RelationCopyStorageUsingBuffer(RelFileLocator srclocator,
 		UnlockReleaseBuffer(dstBuf);
 		UnlockReleaseBuffer(srcBuf);
 	}
-	Assert(read_stream_next_buffer(src_stream, NULL) == InvalidBuffer);
+	Assert(read_stream_get_buffer(src_stream) == InvalidBuffer);
 	read_stream_end(src_stream);
 
 	FreeAccessStrategy(bstrategy_src);

src/include/storage/read_stream.h

@@ -91,6 +91,7 @@ extern ReadStream *read_stream_begin_relation(int flags,
 extern Buffer read_stream_next_buffer(ReadStream *stream, void **per_buffer_data);
 extern BlockNumber read_stream_next_block(ReadStream *stream,
 										  BufferAccessStrategy *strategy);
+extern size_t read_stream_per_buffer_data_size(ReadStream *stream);
 extern ReadStream *read_stream_begin_smgr_relation(int flags,
 												   BufferAccessStrategy strategy,
 												   SMgrRelation smgr,
@@ -102,4 +103,67 @@ extern ReadStream *read_stream_begin_smgr_relation(int flags,
 extern void read_stream_reset(ReadStream *stream);
 extern void read_stream_end(ReadStream *stream);
 
+/*
+ * Get the next buffer from a stream that is not using per-buffer data.
+ */
+static inline Buffer
+read_stream_get_buffer(ReadStream *stream)
+{
+	Assert(read_stream_per_buffer_data_size(stream) == 0);
+	return read_stream_next_buffer(stream, NULL);
+}
+
+/*
+ * Helper for read_stream_get_buffer_and_value().
+ */
+static inline Buffer
+read_stream_get_buffer_and_value_with_size(ReadStream *stream,
+										   void *output_data,
+										   size_t output_data_size)
+{
+	Buffer		buffer;
+	void	   *per_buffer_data;
+
+	Assert(read_stream_per_buffer_data_size(stream) == output_data_size);
+	buffer = read_stream_next_buffer(stream, &per_buffer_data);
+	if (buffer != InvalidBuffer)
+		memcpy(output_data, per_buffer_data, output_data_size);
+
+	return buffer;
+}
+
+/*
+ * Get the next buffer and a copy of the associated per-buffer data.
+ * InvalidBuffer means end-of-stream, and in that case the per-buffer data is
+ * undefined.  Example of use:
+ *
+ * int my_int;
+ *
+ * buf = read_stream_get_buffer_and_value(stream, &my_int);
+ */
+#define read_stream_get_buffer_and_value(stream, vp) \
+	read_stream_get_buffer_and_value_with_size((stream), (vp), sizeof(*(vp)))
+
+/*
+ * Get the next buffer and a pointer to the associated per-buffer data.  This
+ * avoids casts in the calling code, and asserts that we received a pointer to
+ * a pointer to a type that doesn't exceed the storage size.  For example:
+ *
+ * int *my_int_p;
+ *
+ * buf = read_stream_get_buffer_and_pointer(stream, &my_int_p);
+ */
+#define read_stream_get_buffer_and_pointer(stream, pointer) \
+	(AssertMacro(sizeof(**(pointer)) <= read_stream_per_buffer_data_size(stream)), \
+	 read_stream_next_buffer((stream), ((void **) (pointer))))
+
+/*
+ * Set the per-buffer data by value.  This can be called from inside a
+ * callback that is returning block numbers.  It asserts that the value's size
+ * matches the available space.
+ */
+#define read_stream_put_value(stream, per_buffer_data, value) \
+	(AssertMacro(sizeof(value) == read_stream_per_buffer_data_size(stream)), \
+	 memcpy((per_buffer_data), &(value), sizeof(value)))
+
 #endif							/* READ_STREAM_H */