From 5baab6990ee5dd9d828b135f1a28f3eed403ff03 Mon Sep 17 00:00:00 2001
From: John Naylor <john.naylor@postgresql.org>
Date: Fri, 23 Jun 2023 15:41:39 +0700
Subject: [PATCH v35 6/7] Add tidstore tests to benchmark

Extracted from v32 0005 by Masahiko Sawada
---
 .../bench_radix_tree--1.0.sql                 | 10 +++
 contrib/bench_radix_tree/bench_radix_tree.c   | 65 ++++++++++++++++++-
 2 files changed, 72 insertions(+), 3 deletions(-)

diff --git a/contrib/bench_radix_tree/bench_radix_tree--1.0.sql b/contrib/bench_radix_tree/bench_radix_tree--1.0.sql
index db33a1a828..ad66265e23 100644
--- a/contrib/bench_radix_tree/bench_radix_tree--1.0.sql
+++ b/contrib/bench_radix_tree/bench_radix_tree--1.0.sql
@@ -76,3 +76,13 @@ returns record
 as 'MODULE_PATHNAME'
 LANGUAGE C STRICT VOLATILE PARALLEL UNSAFE;
 
+create function bench_tidstore_load(
+minblk int4,
+maxblk int4,
+OUT mem_allocated int8,
+OUT load_ms int8,
+OUT iter_ms int8
+)
+returns record
+as 'MODULE_PATHNAME'
+LANGUAGE C STRICT VOLATILE PARALLEL UNSAFE;
diff --git a/contrib/bench_radix_tree/bench_radix_tree.c b/contrib/bench_radix_tree/bench_radix_tree.c
index 81ada0fd8f..a3aba12aad 100644
--- a/contrib/bench_radix_tree/bench_radix_tree.c
+++ b/contrib/bench_radix_tree/bench_radix_tree.c
@@ -9,6 +9,7 @@
  */
 #include "postgres.h"
 
+#include "access/tidstore.h"
 #include "common/pg_prng.h"
 #include "fmgr.h"
 #include "funcapi.h"
@@ -54,6 +55,7 @@ PG_FUNCTION_INFO_V1(bench_load_random_int);
 PG_FUNCTION_INFO_V1(bench_fixed_height_search);
 PG_FUNCTION_INFO_V1(bench_search_random_nodes);
 PG_FUNCTION_INFO_V1(bench_node128_load);
+PG_FUNCTION_INFO_V1(bench_tidstore_load);
 
 static uint64
 tid_to_key_off(ItemPointer tid, uint32 *off)
@@ -168,6 +170,64 @@ vac_cmp_itemptr(const void *left, const void *right)
 }
 #endif
 
+Datum
+bench_tidstore_load(PG_FUNCTION_ARGS)
+{
+	BlockNumber minblk = PG_GETARG_INT32(0);
+	BlockNumber maxblk = PG_GETARG_INT32(1);
+	TidStore	*ts;
+	TidStoreIter *iter;
+	TidStoreIterResult *result;
+	OffsetNumber *offs;
+	TimestampTz start_time,
+				end_time;
+	long		secs;
+	int			usecs;
+	int64		load_ms;
+	int64		iter_ms;
+	TupleDesc	tupdesc;
+	Datum		values[3];
+	bool		nulls[3] = {false};
+
+	/* Build a tuple descriptor for our result type */
+	if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
+		elog(ERROR, "return type must be a row type");
+
+	offs = palloc(sizeof(OffsetNumber) * TIDS_PER_BLOCK_FOR_LOAD);
+	for (int i = 0; i < TIDS_PER_BLOCK_FOR_LOAD; i++)
+		offs[i] = i + 1; /* FirstOffsetNumber is 1 */
+
+	ts = TidStoreCreate(1 * 1024L * 1024L * 1024L, MaxHeapTuplesPerPage, NULL);
+
+	/* load tids */
+	start_time = GetCurrentTimestamp();
+	for (BlockNumber blkno = minblk; blkno < maxblk; blkno++)
+		TidStoreSetBlockOffsets(ts, blkno, offs, TIDS_PER_BLOCK_FOR_LOAD);
+	end_time = GetCurrentTimestamp();
+	TimestampDifference(start_time, end_time, &secs, &usecs);
+	load_ms = secs * 1000 + usecs / 1000;
+
+	elog(NOTICE, "sleeping for 2 seconds...");
+	pg_usleep(2 * 1000000L);
+
+	/* iterate through tids */
+	iter = TidStoreBeginIterate(ts);
+	start_time = GetCurrentTimestamp();
+	while ((result = TidStoreIterateNext(iter)) != NULL)
+		;
+	TidStoreEndIterate(iter);
+	end_time = GetCurrentTimestamp();
+	TimestampDifference(start_time, end_time, &secs, &usecs);
+	iter_ms = secs * 1000 + usecs / 1000;
+
+	values[0] = Int64GetDatum(TidStoreMemoryUsage(ts));
+	values[1] = Int64GetDatum(load_ms);
+	values[2] = Int64GetDatum(iter_ms);
+
+	TidStoreDestroy(ts);
+	PG_RETURN_DATUM(HeapTupleGetDatum(heap_form_tuple(tupdesc, values, nulls)));
+}
+
 static Datum
 bench_search(FunctionCallInfo fcinfo, bool shuffle)
 {
@@ -663,7 +723,7 @@ bench_node128_load(PG_FUNCTION_ARGS)
 	rt_free(rt);
 	PG_RETURN_DATUM(HeapTupleGetDatum(heap_form_tuple(tupdesc, values, nulls)));
 }
-#if 1
+
 /* to silence warnings about unused iter functions */
 static void pg_attribute_unused()
 stub_iter()
@@ -678,5 +738,4 @@ stub_iter()
 	iter = rt_begin_iterate(rt);
 	rt_iterate_next(iter, &key, &value);
 	rt_end_iterate(iter);
-}
-#endif
+}
\ No newline at end of file
-- 
2.41.0