From f819ff0228c85151e798d861b376e99e571f0cfe Mon Sep 17 00:00:00 2001 From: Thomas Munro Date: Fri, 28 Feb 2025 10:48:29 +1300 Subject: [PATCH 1/2] 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. --- contrib/pg_prewarm/pg_prewarm.c | 4 +- contrib/pg_visibility/pg_visibility.c | 6 +-- src/backend/access/heap/heapam.c | 2 +- src/backend/access/heap/heapam_handler.c | 2 +- src/backend/access/heap/vacuumlazy.c | 6 +-- src/backend/storage/aio/read_stream.c | 12 ++++++ src/backend/storage/buffer/bufmgr.c | 4 +- src/include/storage/read_stream.h | 55 ++++++++++++++++++++++++ 8 files changed, 78 insertions(+), 13 deletions(-) diff --git a/contrib/pg_prewarm/pg_prewarm.c b/contrib/pg_prewarm/pg_prewarm.c index c0efb530c4e5..20bdc3d98660 100644 --- a/contrib/pg_prewarm/pg_prewarm.c +++ b/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); } diff --git a/contrib/pg_visibility/pg_visibility.c b/contrib/pg_visibility/pg_visibility.c index d79ef35006bf..1b8a375778c7 100644 --- a/contrib/pg_visibility/pg_visibility.c +++ b/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; diff --git a/src/backend/access/heap/heapam.c b/src/backend/access/heap/heapam.c index ed2e30217992..7a0c1f21516a 100644 --- a/src/backend/access/heap/heapam.c +++ b/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); } diff --git a/src/backend/access/heap/heapam_handler.c b/src/backend/access/heap/heapam_handler.c index ac082fefa77a..53d7b534c104 100644 --- a/src/backend/access/heap/heapam_handler.c +++ b/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; diff --git a/src/backend/access/heap/vacuumlazy.c b/src/backend/access/heap/vacuumlazy.c index f28326bad095..082d2e067562 100644 --- a/src/backend/access/heap/vacuumlazy.c +++ b/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); @@ -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)) diff --git a/src/backend/storage/aio/read_stream.c b/src/backend/storage/aio/read_stream.c index 36c54fb695b0..a28f7f20b6a0 100644 --- a/src/backend/storage/aio/read_stream.c +++ b/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 diff --git a/src/backend/storage/buffer/bufmgr.c b/src/backend/storage/buffer/bufmgr.c index 1c37d7dfe2f9..b6cc5eea2db7 100644 --- a/src/backend/storage/buffer/bufmgr.c +++ b/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); diff --git a/src/include/storage/read_stream.h b/src/include/storage/read_stream.h index 9b0d65161d02..8aaf3eb3292e 100644 --- a/src/include/storage/read_stream.h +++ b/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 */ From c4807249263271a750952302fdcc038287d9b06c Mon Sep 17 00:00:00 2001 From: Thomas Munro Date: Fri, 28 Feb 2025 12:53:23 +1300 Subject: [PATCH 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 082d2e067562..19fb43e94fa1 100644 --- a/src/backend/access/heap/vacuumlazy.c +++ b/src/backend/access/heap/vacuumlazy.c @@ -1630,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 @@ -1646,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; } } @@ -2691,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; } diff --git a/src/include/storage/read_stream.h b/src/include/storage/read_stream.h index 8aaf3eb3292e..1956663ba86f 100644 --- a/src/include/storage/read_stream.h +++ b/src/include/storage/read_stream.h @@ -157,4 +157,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 */