From 15e740bf3f7a58fccce3f19a94c840ff85016049 Mon Sep 17 00:00:00 2001
From: John Naylor <john.naylor@postgresql.org>
Date: Wed, 13 Mar 2024 14:02:01 +0700
Subject: [PATCH v71 4/6] WIP: Use array of TIDs in TID store regression tests

---
 .../test_tidstore/expected/test_tidstore.out  |  53 ++-----
 .../test_tidstore/sql/test_tidstore.sql       |   7 +-
 .../test_tidstore/test_tidstore--1.0.sql      |   7 +-
 .../modules/test_tidstore/test_tidstore.c     | 135 ++++++++++++------
 4 files changed, 114 insertions(+), 88 deletions(-)

diff --git a/src/test/modules/test_tidstore/expected/test_tidstore.out b/src/test/modules/test_tidstore/expected/test_tidstore.out
index 80698924a1..4eab5d30ba 100644
--- a/src/test/modules/test_tidstore/expected/test_tidstore.out
+++ b/src/test/modules/test_tidstore/expected/test_tidstore.out
@@ -53,16 +53,16 @@ VALUES (0), (1), (:maxblkno - 1), (:maxblkno / 2), (:maxblkno)
 offsets (off) AS (
 VALUES (1), (2), (:maxoffset / 2), (:maxoffset - 1), (:maxoffset)
 )
-SELECT test_set_block_offsets(blk, array_agg(offsets.off)::int2[])
+SELECT do_set_block_offsets(blk, array_agg(offsets.off)::int2[])
   FROM blocks, offsets
   GROUP BY blk;
- test_set_block_offsets 
-------------------------
-             2147483647
-                      0
-             4294967294
-                      1
-             4294967295
+ do_set_block_offsets 
+----------------------
+           2147483647
+                    0
+           4294967294
+                    1
+           4294967295
 (5 rows)
 
 -- Lookup test and dump (sorted) tids.
@@ -81,42 +81,19 @@ SELECT lookup_test();
  (4294967295,512)
 (10 rows)
 
+-- Check TIDs we've added to the store.
+SELECT check_set_block_offsets();
+ check_set_block_offsets 
+-------------------------
+ 
+(1 row)
+
 SELECT test_is_full();
  test_is_full 
 --------------
  f
 (1 row)
 
-SELECT test_dump_tids();
-  test_dump_tids  
-------------------
- (0,1)
- (0,2)
- (0,256)
- (0,511)
- (0,512)
- (1,1)
- (1,2)
- (1,256)
- (1,511)
- (1,512)
- (2147483647,1)
- (2147483647,2)
- (2147483647,256)
- (2147483647,511)
- (2147483647,512)
- (4294967294,1)
- (4294967294,2)
- (4294967294,256)
- (4294967294,511)
- (4294967294,512)
- (4294967295,1)
- (4294967295,2)
- (4294967295,256)
- (4294967295,511)
- (4294967295,512)
-(25 rows)
-
 -- cleanup
 SELECT test_destroy();
  test_destroy 
diff --git a/src/test/modules/test_tidstore/sql/test_tidstore.sql b/src/test/modules/test_tidstore/sql/test_tidstore.sql
index cc1207821c..7317f52058 100644
--- a/src/test/modules/test_tidstore/sql/test_tidstore.sql
+++ b/src/test/modules/test_tidstore/sql/test_tidstore.sql
@@ -43,14 +43,17 @@ VALUES (0), (1), (:maxblkno - 1), (:maxblkno / 2), (:maxblkno)
 offsets (off) AS (
 VALUES (1), (2), (:maxoffset / 2), (:maxoffset - 1), (:maxoffset)
 )
-SELECT test_set_block_offsets(blk, array_agg(offsets.off)::int2[])
+SELECT do_set_block_offsets(blk, array_agg(offsets.off)::int2[])
   FROM blocks, offsets
   GROUP BY blk;
 
 -- Lookup test and dump (sorted) tids.
 SELECT lookup_test();
+
+-- Check TIDs we've added to the store.
+SELECT check_set_block_offsets();
+
 SELECT test_is_full();
-SELECT test_dump_tids();
 
 -- cleanup
 SELECT test_destroy();
diff --git a/src/test/modules/test_tidstore/test_tidstore--1.0.sql b/src/test/modules/test_tidstore/test_tidstore--1.0.sql
index 305459334d..c9f8dfd6f9 100644
--- a/src/test/modules/test_tidstore/test_tidstore--1.0.sql
+++ b/src/test/modules/test_tidstore/test_tidstore--1.0.sql
@@ -8,15 +8,14 @@ shared bool)
 RETURNS void STRICT PARALLEL UNSAFE
 AS 'MODULE_PATHNAME' LANGUAGE C;
 
-CREATE FUNCTION test_set_block_offsets(
+CREATE FUNCTION do_set_block_offsets(
 blkno bigint,
 offsets int2[])
 RETURNS bigint STRICT PARALLEL UNSAFE
 AS 'MODULE_PATHNAME' LANGUAGE C;
 
-CREATE FUNCTION test_dump_tids(
-t_ctid OUT tid)
-RETURNS SETOF tid STRICT PARALLEL UNSAFE
+CREATE FUNCTION check_set_block_offsets()
+RETURNS void STRICT PARALLEL UNSAFE
 AS 'MODULE_PATHNAME' LANGUAGE C;
 
 CREATE FUNCTION test_lookup_tids(
diff --git a/src/test/modules/test_tidstore/test_tidstore.c b/src/test/modules/test_tidstore/test_tidstore.c
index 428d6a3fcf..f5b65eace0 100644
--- a/src/test/modules/test_tidstore/test_tidstore.c
+++ b/src/test/modules/test_tidstore/test_tidstore.c
@@ -24,17 +24,54 @@
 PG_MODULE_MAGIC;
 
 PG_FUNCTION_INFO_V1(test_create);
-PG_FUNCTION_INFO_V1(test_set_block_offsets);
-PG_FUNCTION_INFO_V1(test_dump_tids);
+PG_FUNCTION_INFO_V1(do_set_block_offsets);
+PG_FUNCTION_INFO_V1(check_set_block_offsets);
 PG_FUNCTION_INFO_V1(test_lookup_tids);
 PG_FUNCTION_INFO_V1(test_is_full);
 PG_FUNCTION_INFO_V1(test_destroy);
 
 static TidStore *tidstore = NULL;
 static dsa_area *dsa = NULL;
-static int64 num_tids = 0;
 static size_t max_bytes = (2 * 1024 * 1024L);	/* 2MB */
 
+/* array for verification of some tests */
+typedef struct ItemArray
+{
+	ItemPointerData	*tids;
+	int		max_tids;
+	int		num_tids;
+} ItemArray;
+
+static ItemArray items;
+
+/* comparator routine for ItemPointer */
+static int
+itemptr_cmp(const void *left, const void *right)
+{
+	BlockNumber lblk,
+		rblk;
+	OffsetNumber loff,
+		roff;
+
+	lblk = ItemPointerGetBlockNumber((ItemPointer) left);
+	rblk = ItemPointerGetBlockNumber((ItemPointer) right);
+
+	if (lblk < rblk)
+		return -1;
+	if (lblk > rblk)
+		return 1;
+
+	loff = ItemPointerGetOffsetNumber((ItemPointer) left);
+	roff = ItemPointerGetOffsetNumber((ItemPointer) right);
+
+	if (loff < roff)
+		return -1;
+	if (loff > roff)
+		return 1;
+
+	return 0;
+}
+
 /*
  * Create a TidStore. If shared is false, the tidstore is created
  * on TopMemoryContext, otherwise on DSA. Although the tidstore
@@ -72,7 +109,9 @@ test_create(PG_FUNCTION_ARGS)
 	else
 		tidstore = TidStoreCreate(max_bytes, NULL);
 
-	num_tids = 0;
+	items.num_tids = 0;
+	items.max_tids = max_bytes / sizeof(ItemPointerData);
+	items.tids =  (ItemPointerData *) palloc0(max_bytes);
 
 	MemoryContextSwitchTo(old_ctx);
 
@@ -95,7 +134,7 @@ sanity_check_array(ArrayType *ta)
 
 /* Set the given block and offsets pairs */
 Datum
-test_set_block_offsets(PG_FUNCTION_ARGS)
+do_set_block_offsets(PG_FUNCTION_ARGS)
 {
 	BlockNumber blkno = PG_GETARG_INT64(0);
 	ArrayType  *ta = PG_GETARG_ARRAYTYPE_P_COPY(1);
@@ -107,39 +146,58 @@ test_set_block_offsets(PG_FUNCTION_ARGS)
 	noffs = ArrayGetNItems(ARR_NDIM(ta), ARR_DIMS(ta));
 	offs = ((OffsetNumber *) ARR_DATA_PTR(ta));
 
-	/* Set TIDs */
+	/* Set TIDs in the store */
 	TidStoreLockExclusive(tidstore);
 	TidStoreSetBlockOffsets(tidstore, blkno, offs, noffs);
 	TidStoreUnlock(tidstore);
 
+	/* Set TIDs in verification array */
+	for (int i = 0; i < noffs; i++)
+	{
+		ItemPointer tid;
+		int idx = items.num_tids + i;
+
+		/* Enlarge the TID array if necessary */
+		if (idx >= items.max_tids)
+		{
+			items.max_tids *= 2;
+			items.tids = repalloc(items.tids, sizeof(ItemPointerData) * items.max_tids);
+		}
+
+		tid = &(items.tids[idx]);
+
+		ItemPointerSetBlockNumber(tid, blkno);
+		ItemPointerSetOffsetNumber(tid, offs[i]);
+	}
+
 	/* Update statistics */
-	num_tids += noffs;
+	items.num_tids += noffs;
 
 	PG_RETURN_INT64(blkno);
 }
 
 /*
- * Dump and return TIDs in the tidstore. The output TIDs are ordered.
+ * Verify TIDs in store against the array.
  */
 Datum
-test_dump_tids(PG_FUNCTION_ARGS)
+check_set_block_offsets(PG_FUNCTION_ARGS)
 {
-	FuncCallContext *funcctx;
-	ItemPointerData *tids;
-
-	if (SRF_IS_FIRSTCALL())
-	{
-		MemoryContext oldcontext;
 		TidStoreIter *iter;
 		TidStoreIterResult *iter_result;
-		int64		ntids = 0;
+		ItemPointerData * tids;
+		int ntids = 0;
+
+		/* lookup each member in the verification array */
+		for (int i = 0; i < items.num_tids; i++)
+			if (!TidStoreIsMember(tidstore, &items.tids[i]))
+				elog(ERROR, "missing TID with block %u, offset %u",
+					ItemPointerGetBlockNumber(&items.tids[i]),
+					ItemPointerGetOffsetNumber(&items.tids[i]));
 
-		funcctx = SRF_FIRSTCALL_INIT();
-		oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
 		tids = (ItemPointerData *)
-			palloc0(sizeof(ItemPointerData) * num_tids);
+			palloc0(sizeof(ItemPointerData) * items.num_tids);
 
-		/* Collect TIDs stored in the tidstore */
+		/* Collect TIDs stored in the tidstore, in order */
 		TidStoreLockShare(tidstore);
 		iter = TidStoreBeginIterate(tidstore);
 		while ((iter_result = TidStoreIterateNext(iter)) != NULL)
@@ -148,33 +206,21 @@ test_dump_tids(PG_FUNCTION_ARGS)
 				ItemPointerSet(&(tids[ntids++]), iter_result->blkno,
 							   iter_result->offsets[i]);
 		}
+		TidStoreEndIterate(iter);
 		TidStoreUnlock(tidstore);
 
-		Assert(ntids == num_tids);
-
-		funcctx->user_fctx = tids;
-		funcctx->max_calls = num_tids;
-
-		MemoryContextSwitchTo(oldcontext);
-	}
-
-	funcctx = SRF_PERCALL_SETUP();
-	tids = (ItemPointerData *) funcctx->user_fctx;
+		if (ntids != items.num_tids)
+			elog(ERROR, "should have %d TIDs, have %d", items.num_tids, ntids);
 
-	if (funcctx->call_cntr < funcctx->max_calls)
-	{
-		int			idx;
-
-		/*
-		 * Note that since funcctx->call_cntr is incremented in
-		 * SRF_RETURN_NEXT before return, we need to remember the current
-		 * counter to access the tid array.
-		 */
-		idx = funcctx->call_cntr;
-		SRF_RETURN_NEXT(funcctx, PointerGetDatum(&(tids[idx])));
-	}
+		/* Sort verification array and compare each member to what we dumped from the TID store. */
+		qsort(items.tids, items.num_tids, sizeof(ItemPointerData), itemptr_cmp);
+		for (int i = 0; i < items.num_tids; i++)
+		{
+			if (itemptr_cmp((const void *) &items.tids[i], (const void *) &tids[i]) != 0)
+				elog(ERROR, "Dumped TID array doesn't match verification array");
+		}
 
-	SRF_RETURN_DONE(funcctx);
+	PG_RETURN_VOID();
 }
 
 /*
@@ -236,7 +282,8 @@ test_destroy(PG_FUNCTION_ARGS)
 {
 	TidStoreDestroy(tidstore);
 	tidstore = NULL;
-	num_tids = 0;
+	items.num_tids = 0;
+	pfree(items.tids);
 
 	if (dsa)
 		dsa_detach(dsa);
-- 
2.44.0

