From 2586d9c7321391168a40cb0f14e5a80182792b64 Mon Sep 17 00:00:00 2001
From: Thomas Munro <thomas.munro@gmail.com>
Date: Fri, 28 Feb 2025 12:53:23 +1300
Subject: [PATCH v2 2/2] Improve API for storing data in read streams.

Read stream callbacks receive a void pointer into the per-buffer data
queue so that can store data there for later retrieval by the buffer
consumer.  We can improve readability and safety a bit by changing
cast-and-assign or raw memcpy() to:

	read_stream_put_value(stream, per_buffer_data, my_int);

This form infers the size and asserts that the storage space matches,
generally mirroring the read_stream_get_buffer_and_value() call used for
retrieving the streamed data later.
---
 src/backend/access/heap/vacuumlazy.c | 6 +++---
 src/include/storage/read_stream.h    | 9 +++++++++
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/src/backend/access/heap/vacuumlazy.c b/src/backend/access/heap/vacuumlazy.c
index ac7a4d8c21d..9563906fb27 100644
--- a/src/backend/access/heap/vacuumlazy.c
+++ b/src/backend/access/heap/vacuumlazy.c
@@ -1612,7 +1612,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
@@ -1628,7 +1628,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;
 	}
 }
@@ -2671,7 +2671,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;
 }
diff --git a/src/include/storage/read_stream.h b/src/include/storage/read_stream.h
index c6066c0f296..5af801969b4 100644
--- a/src/include/storage/read_stream.h
+++ b/src/include/storage/read_stream.h
@@ -136,4 +136,13 @@ read_stream_get_buffer_and_value_with_size(ReadStream *stream,
 	(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 */
-- 
2.48.1

