Improve API for retrieving data from read streams.

Dealing with the per_buffer_data argument to read_stream_next_buffer()
has proven a bit clunky.  Provide some new wrapper functions/macros:

    buffer = read_stream_get_buffer(rs);
    buffer = read_stream_get_buffer_and_value(rs, &my_int);
    buffer = read_stream_get_buffer_and_pointer(rs, &my_pointer_to_int);

These improve readability and type safety via assertions.

Author: macdice

contrib/pg_prewarm/pg_prewarm.c

@@ -225,11 +225,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

@@ -672,7 +672,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

@@ -1247,7 +1247,6 @@ lazy_scan_heap(LVRelState *vacrel)
 		uint8		blk_info = 0;
 		int			ndeleted = 0;
 		bool		has_lpdead_items;
-		void	   *per_buffer_data = NULL;
 		bool		vm_page_frozen = false;
 		bool		got_cleanup_lock = false;
 
@@ -1307,13 +1306,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);
@@ -2769,7 +2767,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

@@ -786,6 +786,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)
@@ -1034,6 +1037,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

@@ -5188,7 +5188,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);
 
@@ -5213,7 +5213,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,58 @@ 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))))
+
 #endif							/* READ_STREAM_H */