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-2025, 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 2 : PG_MODULE_MAGIC;
28 :
29 4 : PG_FUNCTION_INFO_V1(test_create);
30 4 : PG_FUNCTION_INFO_V1(do_set_block_offsets);
31 4 : PG_FUNCTION_INFO_V1(check_set_block_offsets);
32 4 : PG_FUNCTION_INFO_V1(test_is_full);
33 4 : 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 415912 : itemptr_cmp(const void *left, const void *right)
53 : {
54 : BlockNumber lblk,
55 : rblk;
56 : OffsetNumber loff,
57 : roff;
58 :
59 415912 : lblk = ItemPointerGetBlockNumber((ItemPointer) left);
60 415912 : rblk = ItemPointerGetBlockNumber((ItemPointer) right);
61 :
62 415912 : if (lblk < rblk)
63 128168 : return -1;
64 287744 : if (lblk > rblk)
65 129020 : return 1;
66 :
67 158724 : loff = ItemPointerGetOffsetNumber((ItemPointer) left);
68 158724 : roff = ItemPointerGetOffsetNumber((ItemPointer) right);
69 :
70 158724 : if (loff < roff)
71 70128 : return -1;
72 88596 : if (loff > roff)
73 27684 : return 1;
74 :
75 60912 : return 0;
76 : }
77 :
78 : /*
79 : * Create a TidStore. If shared is false, the tidstore is created
80 : * on TopMemoryContext, otherwise on DSA. Although the tidstore
81 : * is created on DSA, only the same process can subsequently use
82 : * the tidstore. The tidstore handle is not shared anywhere.
83 : */
84 : Datum
85 6 : test_create(PG_FUNCTION_ARGS)
86 : {
87 6 : bool shared = PG_GETARG_BOOL(0);
88 : MemoryContext old_ctx;
89 :
90 : /* doesn't really matter, since it's just a hint */
91 6 : size_t tidstore_max_size = 2 * 1024 * 1024;
92 6 : size_t array_init_size = 1024;
93 :
94 : Assert(tidstore == NULL);
95 :
96 : /*
97 : * Create the TidStore on TopMemoryContext so that the same process use it
98 : * for subsequent tests.
99 : */
100 6 : old_ctx = MemoryContextSwitchTo(TopMemoryContext);
101 :
102 6 : if (shared)
103 : {
104 : int tranche_id;
105 :
106 2 : tranche_id = LWLockNewTrancheId("test_tidstore");
107 :
108 2 : tidstore = TidStoreCreateShared(tidstore_max_size, tranche_id);
109 :
110 : /*
111 : * Remain attached until end of backend or explicitly detached so that
112 : * the same process use the tidstore for subsequent tests.
113 : */
114 2 : dsa_pin_mapping(TidStoreGetDSA(tidstore));
115 : }
116 : else
117 : /* VACUUM uses insert only, so we test the other option. */
118 4 : tidstore = TidStoreCreateLocal(tidstore_max_size, false);
119 :
120 6 : tidstore_empty_size = TidStoreMemoryUsage(tidstore);
121 :
122 6 : items.num_tids = 0;
123 6 : items.max_tids = array_init_size / sizeof(ItemPointerData);
124 6 : items.insert_tids = (ItemPointerData *) palloc0(array_init_size);
125 6 : items.lookup_tids = (ItemPointerData *) palloc0(array_init_size);
126 6 : items.iter_tids = (ItemPointerData *) palloc0(array_init_size);
127 :
128 6 : MemoryContextSwitchTo(old_ctx);
129 :
130 6 : PG_RETURN_VOID();
131 : }
132 :
133 : static void
134 2242 : sanity_check_array(ArrayType *ta)
135 : {
136 2242 : if (ARR_HASNULL(ta) && array_contains_nulls(ta))
137 0 : ereport(ERROR,
138 : (errcode(ERRCODE_NULL_VALUE_NOT_ALLOWED),
139 : errmsg("array must not contain nulls")));
140 :
141 2242 : if (ARR_NDIM(ta) > 1)
142 0 : ereport(ERROR,
143 : (errcode(ERRCODE_DATA_EXCEPTION),
144 : errmsg("argument must be empty or one-dimensional array")));
145 2242 : }
146 :
147 : static void
148 2276 : check_tidstore_available(void)
149 : {
150 2276 : if (tidstore == NULL)
151 0 : elog(ERROR, "tidstore is not created");
152 2276 : }
153 :
154 : static void
155 2240 : purge_from_verification_array(BlockNumber blkno)
156 : {
157 2240 : int dst = 0;
158 :
159 6982660 : for (int src = 0; src < items.num_tids; src++)
160 6980420 : if (ItemPointerGetBlockNumber(&items.insert_tids[src]) != blkno)
161 6980372 : items.insert_tids[dst++] = items.insert_tids[src];
162 2240 : items.num_tids = dst;
163 2240 : }
164 :
165 :
166 : /* Set the given block and offsets pairs */
167 : Datum
168 2242 : do_set_block_offsets(PG_FUNCTION_ARGS)
169 : {
170 2242 : BlockNumber blkno = PG_GETARG_INT64(0);
171 2242 : ArrayType *ta = PG_GETARG_ARRAYTYPE_P_COPY(1);
172 : OffsetNumber *offs;
173 : int noffs;
174 :
175 2242 : check_tidstore_available();
176 2242 : sanity_check_array(ta);
177 :
178 2242 : noffs = ArrayGetNItems(ARR_NDIM(ta), ARR_DIMS(ta));
179 2242 : offs = ((OffsetNumber *) ARR_DATA_PTR(ta));
180 :
181 : /* Set TIDs in the store */
182 2242 : TidStoreLockExclusive(tidstore);
183 2242 : TidStoreSetBlockOffsets(tidstore, blkno, offs, noffs);
184 2240 : TidStoreUnlock(tidstore);
185 :
186 : /* Remove the existing items of blkno from the verification array */
187 2240 : purge_from_verification_array(blkno);
188 :
189 : /* Set TIDs in verification array */
190 32696 : for (int i = 0; i < noffs; i++)
191 : {
192 : ItemPointer tid;
193 30456 : int idx = items.num_tids + i;
194 :
195 : /* Enlarge the TID arrays if necessary */
196 30456 : if (idx >= items.max_tids)
197 : {
198 24 : items.max_tids *= 2;
199 24 : items.insert_tids = repalloc(items.insert_tids, sizeof(ItemPointerData) * items.max_tids);
200 24 : items.lookup_tids = repalloc(items.lookup_tids, sizeof(ItemPointerData) * items.max_tids);
201 24 : items.iter_tids = repalloc(items.iter_tids, sizeof(ItemPointerData) * items.max_tids);
202 : }
203 :
204 30456 : tid = &(items.insert_tids[idx]);
205 30456 : ItemPointerSet(tid, blkno, offs[i]);
206 : }
207 :
208 : /* Update statistics */
209 2240 : items.num_tids += noffs;
210 :
211 2240 : PG_RETURN_INT64(blkno);
212 : }
213 :
214 : /*
215 : * Verify TIDs in store against the array.
216 : */
217 : Datum
218 24 : check_set_block_offsets(PG_FUNCTION_ARGS)
219 : {
220 : TidStoreIter *iter;
221 : TidStoreIterResult *iter_result;
222 24 : int num_iter_tids = 0;
223 24 : int num_lookup_tids = 0;
224 24 : BlockNumber prevblkno = 0;
225 :
226 24 : check_tidstore_available();
227 :
228 : /* lookup each member in the verification array */
229 30480 : for (int i = 0; i < items.num_tids; i++)
230 30456 : if (!TidStoreIsMember(tidstore, &items.insert_tids[i]))
231 0 : elog(ERROR, "missing TID with block %u, offset %u",
232 : ItemPointerGetBlockNumber(&items.insert_tids[i]),
233 : ItemPointerGetOffsetNumber(&items.insert_tids[i]));
234 :
235 : /*
236 : * Lookup all possible TIDs for each distinct block in the verification
237 : * array and save successful lookups in the lookup array.
238 : */
239 :
240 30480 : for (int i = 0; i < items.num_tids; i++)
241 : {
242 30456 : BlockNumber blkno = ItemPointerGetBlockNumber(&items.insert_tids[i]);
243 :
244 30456 : if (i > 0 && blkno == prevblkno)
245 28216 : continue;
246 :
247 4587520 : for (OffsetNumber offset = FirstOffsetNumber; offset < MaxOffsetNumber; offset++)
248 : {
249 : ItemPointerData tid;
250 :
251 4585280 : ItemPointerSet(&tid, blkno, offset);
252 :
253 4585280 : TidStoreLockShare(tidstore);
254 4585280 : if (TidStoreIsMember(tidstore, &tid))
255 30456 : ItemPointerSet(&items.lookup_tids[num_lookup_tids++], blkno, offset);
256 4585280 : TidStoreUnlock(tidstore);
257 : }
258 :
259 2240 : prevblkno = blkno;
260 : }
261 :
262 : /* Collect TIDs stored in the tidstore, in order */
263 :
264 24 : TidStoreLockShare(tidstore);
265 24 : iter = TidStoreBeginIterate(tidstore);
266 2264 : while ((iter_result = TidStoreIterateNext(iter)) != NULL)
267 : {
268 : OffsetNumber offsets[MaxOffsetNumber];
269 : int num_offsets;
270 :
271 2240 : num_offsets = TidStoreGetBlockOffsets(iter_result, offsets, lengthof(offsets));
272 : Assert(num_offsets <= lengthof(offsets));
273 32696 : for (int i = 0; i < num_offsets; i++)
274 30456 : ItemPointerSet(&(items.iter_tids[num_iter_tids++]), iter_result->blkno,
275 30456 : offsets[i]);
276 : }
277 24 : TidStoreEndIterate(iter);
278 24 : TidStoreUnlock(tidstore);
279 :
280 : /*
281 : * Sort verification and lookup arrays and test that all arrays are the
282 : * same.
283 : */
284 :
285 24 : if (num_lookup_tids != items.num_tids)
286 0 : elog(ERROR, "should have %d TIDs, have %d", items.num_tids, num_lookup_tids);
287 24 : if (num_iter_tids != items.num_tids)
288 0 : elog(ERROR, "should have %d TIDs, have %d", items.num_tids, num_iter_tids);
289 :
290 24 : qsort(items.insert_tids, items.num_tids, sizeof(ItemPointerData), itemptr_cmp);
291 24 : qsort(items.lookup_tids, items.num_tids, sizeof(ItemPointerData), itemptr_cmp);
292 30480 : for (int i = 0; i < items.num_tids; i++)
293 : {
294 30456 : if (itemptr_cmp(&items.insert_tids[i], &items.iter_tids[i]) != 0)
295 0 : elog(ERROR, "TID iter array doesn't match verification array, got (%u,%u) expected (%u,%u)",
296 : ItemPointerGetBlockNumber(&items.iter_tids[i]),
297 : ItemPointerGetOffsetNumber(&items.iter_tids[i]),
298 : ItemPointerGetBlockNumber(&items.insert_tids[i]),
299 : ItemPointerGetOffsetNumber(&items.insert_tids[i]));
300 30456 : if (itemptr_cmp(&items.insert_tids[i], &items.lookup_tids[i]) != 0)
301 0 : elog(ERROR, "TID lookup array doesn't match verification array, got (%u,%u) expected (%u,%u)",
302 : ItemPointerGetBlockNumber(&items.lookup_tids[i]),
303 : ItemPointerGetOffsetNumber(&items.lookup_tids[i]),
304 : ItemPointerGetBlockNumber(&items.insert_tids[i]),
305 : ItemPointerGetOffsetNumber(&items.insert_tids[i]));
306 : }
307 :
308 24 : PG_RETURN_VOID();
309 : }
310 :
311 : /*
312 : * In real world use, we care if the memory usage is greater than
313 : * some configured limit. Here we just want to verify that
314 : * TidStoreMemoryUsage is not broken.
315 : */
316 : Datum
317 4 : test_is_full(PG_FUNCTION_ARGS)
318 : {
319 : bool is_full;
320 :
321 4 : check_tidstore_available();
322 :
323 4 : is_full = (TidStoreMemoryUsage(tidstore) > tidstore_empty_size);
324 :
325 4 : PG_RETURN_BOOL(is_full);
326 : }
327 :
328 : /* Free the tidstore */
329 : Datum
330 6 : test_destroy(PG_FUNCTION_ARGS)
331 : {
332 6 : check_tidstore_available();
333 :
334 6 : TidStoreDestroy(tidstore);
335 6 : tidstore = NULL;
336 6 : items.num_tids = 0;
337 6 : pfree(items.insert_tids);
338 6 : pfree(items.lookup_tids);
339 6 : pfree(items.iter_tids);
340 :
341 6 : PG_RETURN_VOID();
342 : }
|