Line data Source code
1 : /*--------------------------------------------------------------------------
2 : *
3 : * test_tidstore.c
4 : * Test TidStore data structure.
5 : *
6 : * Note: all locking in this test module is useless since there is only
7 : * a single process to use the TidStore. It is meant to be an example of
8 : * usage.
9 : *
10 : * Copyright (c) 2024-2026, PostgreSQL Global Development Group
11 : *
12 : * IDENTIFICATION
13 : * src/test/modules/test_tidstore/test_tidstore.c
14 : *
15 : * -------------------------------------------------------------------------
16 : */
17 : #include "postgres.h"
18 :
19 : #include "access/tidstore.h"
20 : #include "fmgr.h"
21 : #include "storage/block.h"
22 : #include "storage/itemptr.h"
23 : #include "storage/lwlock.h"
24 : #include "utils/array.h"
25 : #include "utils/memutils.h"
26 :
27 1 : PG_MODULE_MAGIC;
28 :
29 2 : PG_FUNCTION_INFO_V1(test_create);
30 2 : PG_FUNCTION_INFO_V1(do_set_block_offsets);
31 2 : PG_FUNCTION_INFO_V1(check_set_block_offsets);
32 2 : PG_FUNCTION_INFO_V1(test_is_full);
33 2 : PG_FUNCTION_INFO_V1(test_destroy);
34 :
35 : static TidStore *tidstore = NULL;
36 : static size_t tidstore_empty_size;
37 :
38 : /* array for verification of some tests */
39 : typedef struct ItemArray
40 : {
41 : ItemPointerData *insert_tids;
42 : ItemPointerData *lookup_tids;
43 : ItemPointerData *iter_tids;
44 : int max_tids;
45 : int num_tids;
46 : } ItemArray;
47 :
48 : static ItemArray items;
49 :
50 : /* comparator routine for ItemPointer */
51 : static int
52 208184 : itemptr_cmp(const void *left, const void *right)
53 : {
54 : BlockNumber lblk,
55 : rblk;
56 : OffsetNumber loff,
57 : roff;
58 :
59 208184 : lblk = ItemPointerGetBlockNumber((const ItemPointerData *) left);
60 208184 : rblk = ItemPointerGetBlockNumber((const ItemPointerData *) right);
61 :
62 208184 : if (lblk < rblk)
63 64084 : return -1;
64 144100 : if (lblk > rblk)
65 64510 : return 1;
66 :
67 79590 : loff = ItemPointerGetOffsetNumber((const ItemPointerData *) left);
68 79590 : roff = ItemPointerGetOffsetNumber((const ItemPointerData *) right);
69 :
70 79590 : if (loff < roff)
71 35178 : return -1;
72 44412 : if (loff > roff)
73 13842 : return 1;
74 :
75 30570 : return 0;
76 : }
77 :
78 : static int
79 14165 : offsetnumber_cmp(const void *a, const void *b)
80 : {
81 14165 : OffsetNumber l = *(const OffsetNumber *) a;
82 14165 : OffsetNumber r = *(const OffsetNumber *) b;
83 :
84 14165 : if (l < r)
85 14165 : return -1;
86 0 : else if (l > r)
87 0 : return 1;
88 0 : return 0;
89 : }
90 :
91 : /*
92 : * Create a TidStore. If shared is false, the tidstore is created
93 : * on TopMemoryContext, otherwise on DSA. Although the tidstore
94 : * is created on DSA, only the same process can subsequently use
95 : * the tidstore. The tidstore handle is not shared anywhere.
96 : */
97 : Datum
98 3 : test_create(PG_FUNCTION_ARGS)
99 : {
100 3 : bool shared = PG_GETARG_BOOL(0);
101 : MemoryContext old_ctx;
102 :
103 : /* doesn't really matter, since it's just a hint */
104 3 : size_t tidstore_max_size = 2 * 1024 * 1024;
105 3 : size_t array_init_size = 1024;
106 :
107 : Assert(tidstore == NULL);
108 :
109 : /*
110 : * Create the TidStore on TopMemoryContext so that the same process use it
111 : * for subsequent tests.
112 : */
113 3 : old_ctx = MemoryContextSwitchTo(TopMemoryContext);
114 :
115 3 : if (shared)
116 : {
117 : int tranche_id;
118 :
119 1 : tranche_id = LWLockNewTrancheId("test_tidstore");
120 :
121 1 : tidstore = TidStoreCreateShared(tidstore_max_size, tranche_id);
122 :
123 : /*
124 : * Remain attached until end of backend or explicitly detached so that
125 : * the same process use the tidstore for subsequent tests.
126 : */
127 1 : dsa_pin_mapping(TidStoreGetDSA(tidstore));
128 : }
129 : else
130 : /* VACUUM uses insert only, so we test the other option. */
131 2 : tidstore = TidStoreCreateLocal(tidstore_max_size, false);
132 :
133 3 : tidstore_empty_size = TidStoreMemoryUsage(tidstore);
134 :
135 3 : items.num_tids = 0;
136 3 : items.max_tids = array_init_size / sizeof(ItemPointerData);
137 3 : items.insert_tids = (ItemPointerData *) palloc0(array_init_size);
138 3 : items.lookup_tids = (ItemPointerData *) palloc0(array_init_size);
139 3 : items.iter_tids = (ItemPointerData *) palloc0(array_init_size);
140 :
141 3 : MemoryContextSwitchTo(old_ctx);
142 :
143 3 : PG_RETURN_VOID();
144 : }
145 :
146 : static void
147 1121 : sanity_check_array(ArrayType *ta)
148 : {
149 1121 : if (ARR_HASNULL(ta) && array_contains_nulls(ta))
150 0 : ereport(ERROR,
151 : (errcode(ERRCODE_NULL_VALUE_NOT_ALLOWED),
152 : errmsg("array must not contain nulls")));
153 :
154 1121 : if (ARR_NDIM(ta) > 1)
155 0 : ereport(ERROR,
156 : (errcode(ERRCODE_DATA_EXCEPTION),
157 : errmsg("argument must be empty or one-dimensional array")));
158 1121 : }
159 :
160 : static void
161 1138 : check_tidstore_available(void)
162 : {
163 1138 : if (tidstore == NULL)
164 0 : elog(ERROR, "tidstore is not created");
165 1138 : }
166 :
167 : static void
168 1120 : purge_from_verification_array(BlockNumber blkno)
169 : {
170 1120 : int dst = 0;
171 :
172 3494634 : for (int src = 0; src < items.num_tids; src++)
173 3493514 : if (ItemPointerGetBlockNumber(&items.insert_tids[src]) != blkno)
174 3493490 : items.insert_tids[dst++] = items.insert_tids[src];
175 1120 : items.num_tids = dst;
176 1120 : }
177 :
178 :
179 : /* Set the given block and offsets pairs */
180 : Datum
181 1121 : do_set_block_offsets(PG_FUNCTION_ARGS)
182 : {
183 1121 : BlockNumber blkno = PG_GETARG_INT64(0);
184 1121 : ArrayType *ta = PG_GETARG_ARRAYTYPE_P_COPY(1);
185 : OffsetNumber *offs;
186 : int noffs;
187 :
188 1121 : check_tidstore_available();
189 1121 : sanity_check_array(ta);
190 :
191 1121 : noffs = ArrayGetNItems(ARR_NDIM(ta), ARR_DIMS(ta));
192 1121 : offs = ((OffsetNumber *) ARR_DATA_PTR(ta));
193 :
194 : /* TidStoreSetBlockOffsets() requires offsets to be strictly ascending. */
195 1121 : qsort(offs, noffs, sizeof(OffsetNumber), offsetnumber_cmp);
196 :
197 : /* Set TIDs in the store */
198 1121 : TidStoreLockExclusive(tidstore);
199 1121 : TidStoreSetBlockOffsets(tidstore, blkno, offs, noffs);
200 1120 : TidStoreUnlock(tidstore);
201 :
202 : /* Remove the existing items of blkno from the verification array */
203 1120 : purge_from_verification_array(blkno);
204 :
205 : /* Set TIDs in verification array */
206 16405 : for (int i = 0; i < noffs; i++)
207 : {
208 : ItemPointer tid;
209 15285 : int idx = items.num_tids + i;
210 :
211 : /* Enlarge the TID arrays if necessary */
212 15285 : if (idx >= items.max_tids)
213 : {
214 12 : items.max_tids *= 2;
215 12 : items.insert_tids = repalloc(items.insert_tids, sizeof(ItemPointerData) * items.max_tids);
216 12 : items.lookup_tids = repalloc(items.lookup_tids, sizeof(ItemPointerData) * items.max_tids);
217 12 : items.iter_tids = repalloc(items.iter_tids, sizeof(ItemPointerData) * items.max_tids);
218 : }
219 :
220 15285 : tid = &(items.insert_tids[idx]);
221 15285 : ItemPointerSet(tid, blkno, offs[i]);
222 : }
223 :
224 : /* Update statistics */
225 1120 : items.num_tids += noffs;
226 :
227 1120 : PG_RETURN_INT64(blkno);
228 : }
229 :
230 : /*
231 : * Verify TIDs in store against the array.
232 : */
233 : Datum
234 12 : check_set_block_offsets(PG_FUNCTION_ARGS)
235 : {
236 : TidStoreIter *iter;
237 : TidStoreIterResult *iter_result;
238 12 : int num_iter_tids = 0;
239 12 : int num_lookup_tids = 0;
240 12 : BlockNumber prevblkno = 0;
241 :
242 12 : check_tidstore_available();
243 :
244 : /* lookup each member in the verification array */
245 15297 : for (int i = 0; i < items.num_tids; i++)
246 15285 : if (!TidStoreIsMember(tidstore, &items.insert_tids[i]))
247 0 : elog(ERROR, "missing TID with block %u, offset %u",
248 : ItemPointerGetBlockNumber(&items.insert_tids[i]),
249 : ItemPointerGetOffsetNumber(&items.insert_tids[i]));
250 :
251 : /*
252 : * Lookup all possible TIDs for each distinct block in the verification
253 : * array and save successful lookups in the lookup array.
254 : */
255 :
256 15297 : for (int i = 0; i < items.num_tids; i++)
257 : {
258 15285 : BlockNumber blkno = ItemPointerGetBlockNumber(&items.insert_tids[i]);
259 :
260 15285 : if (i > 0 && blkno == prevblkno)
261 14165 : continue;
262 :
263 2293760 : for (OffsetNumber offset = FirstOffsetNumber; offset < MaxOffsetNumber; offset++)
264 : {
265 : ItemPointerData tid;
266 :
267 2292640 : ItemPointerSet(&tid, blkno, offset);
268 :
269 2292640 : TidStoreLockShare(tidstore);
270 2292640 : if (TidStoreIsMember(tidstore, &tid))
271 15285 : ItemPointerSet(&items.lookup_tids[num_lookup_tids++], blkno, offset);
272 2292640 : TidStoreUnlock(tidstore);
273 : }
274 :
275 1120 : prevblkno = blkno;
276 : }
277 :
278 : /* Collect TIDs stored in the tidstore, in order */
279 :
280 12 : TidStoreLockShare(tidstore);
281 12 : iter = TidStoreBeginIterate(tidstore);
282 1132 : while ((iter_result = TidStoreIterateNext(iter)) != NULL)
283 : {
284 : OffsetNumber offsets[MaxOffsetNumber];
285 : int num_offsets;
286 :
287 1120 : num_offsets = TidStoreGetBlockOffsets(iter_result, offsets, lengthof(offsets));
288 : Assert(num_offsets <= lengthof(offsets));
289 16405 : for (int i = 0; i < num_offsets; i++)
290 15285 : ItemPointerSet(&(items.iter_tids[num_iter_tids++]), iter_result->blkno,
291 15285 : offsets[i]);
292 : }
293 12 : TidStoreEndIterate(iter);
294 12 : TidStoreUnlock(tidstore);
295 :
296 : /*
297 : * Sort verification and lookup arrays and test that all arrays are the
298 : * same.
299 : */
300 :
301 12 : if (num_lookup_tids != items.num_tids)
302 0 : elog(ERROR, "should have %d TIDs, have %d", items.num_tids, num_lookup_tids);
303 12 : if (num_iter_tids != items.num_tids)
304 0 : elog(ERROR, "should have %d TIDs, have %d", items.num_tids, num_iter_tids);
305 :
306 12 : qsort(items.insert_tids, items.num_tids, sizeof(ItemPointerData), itemptr_cmp);
307 12 : qsort(items.lookup_tids, items.num_tids, sizeof(ItemPointerData), itemptr_cmp);
308 15297 : for (int i = 0; i < items.num_tids; i++)
309 : {
310 15285 : if (itemptr_cmp(&items.insert_tids[i], &items.iter_tids[i]) != 0)
311 0 : elog(ERROR, "TID iter array doesn't match verification array, got (%u,%u) expected (%u,%u)",
312 : ItemPointerGetBlockNumber(&items.iter_tids[i]),
313 : ItemPointerGetOffsetNumber(&items.iter_tids[i]),
314 : ItemPointerGetBlockNumber(&items.insert_tids[i]),
315 : ItemPointerGetOffsetNumber(&items.insert_tids[i]));
316 15285 : if (itemptr_cmp(&items.insert_tids[i], &items.lookup_tids[i]) != 0)
317 0 : elog(ERROR, "TID lookup array doesn't match verification array, got (%u,%u) expected (%u,%u)",
318 : ItemPointerGetBlockNumber(&items.lookup_tids[i]),
319 : ItemPointerGetOffsetNumber(&items.lookup_tids[i]),
320 : ItemPointerGetBlockNumber(&items.insert_tids[i]),
321 : ItemPointerGetOffsetNumber(&items.insert_tids[i]));
322 : }
323 :
324 12 : PG_RETURN_VOID();
325 : }
326 :
327 : /*
328 : * In real world use, we care if the memory usage is greater than
329 : * some configured limit. Here we just want to verify that
330 : * TidStoreMemoryUsage is not broken.
331 : */
332 : Datum
333 2 : test_is_full(PG_FUNCTION_ARGS)
334 : {
335 : bool is_full;
336 :
337 2 : check_tidstore_available();
338 :
339 2 : is_full = (TidStoreMemoryUsage(tidstore) > tidstore_empty_size);
340 :
341 2 : PG_RETURN_BOOL(is_full);
342 : }
343 :
344 : /* Free the tidstore */
345 : Datum
346 3 : test_destroy(PG_FUNCTION_ARGS)
347 : {
348 3 : check_tidstore_available();
349 :
350 3 : TidStoreDestroy(tidstore);
351 3 : tidstore = NULL;
352 3 : items.num_tids = 0;
353 3 : pfree(items.insert_tids);
354 3 : pfree(items.lookup_tids);
355 3 : pfree(items.iter_tids);
356 :
357 3 : PG_RETURN_VOID();
358 : }
|