Age Owner Branch data TLA Line data Source code
1 : : /*-------------------------------------------------------------------------
2 : : *
3 : : * test_shmem.c
4 : : * Helpers to test shmem allocation routines
5 : : *
6 : : * Test basic memory allocation in an extension module. One notable feature
7 : : * that is not exercised by any other module in the repository is the
8 : : * allocating (non-DSM) shared memory after postmaster startup.
9 : : *
10 : : * Copyright (c) 2020-2026, PostgreSQL Global Development Group
11 : : *
12 : : * IDENTIFICATION
13 : : * src/test/modules/test_shmem/test_shmem.c
14 : : *
15 : : *-------------------------------------------------------------------------
16 : : */
17 : :
18 : : #include "postgres.h"
19 : :
20 : : #include "fmgr.h"
21 : : #include "miscadmin.h"
22 : : #include "storage/shmem.h"
23 : : #include "utils/builtins.h"
24 : : #include "utils/guc.h"
25 : : #include "utils/injection_point.h"
26 : :
27 : :
189 heikki.linnakangas@i 28 :CBC 16 : PG_MODULE_MAGIC;
29 : :
30 : : typedef struct TestShmemData
31 : : {
32 : : bool initialized;
33 : : int attach_count;
34 : : char dummy_data[FLEXIBLE_ARRAY_MEMBER];
35 : : } TestShmemData;
36 : :
37 : : static TestShmemData *TestShmem;
38 : :
39 : : #define DEFAULT_TEST_AREA_BYTES sizeof(TestShmemData)
40 : : #define MAX_TEST_AREA_BYTES 1000000
41 : :
42 : : static bool attached_or_initialized = false;
43 : : static int test_shmem_area_size = DEFAULT_TEST_AREA_BYTES;
44 : : static bool test_shmem_guc_defined = false;
45 : :
46 : : static void test_shmem_request(void *arg);
47 : : static void test_shmem_init(void *arg);
48 : : static void test_shmem_attach(void *arg);
49 : :
50 : : static const ShmemCallbacks TestShmemCallbacks = {
51 : : .flags = SHMEM_CALLBACKS_ALLOW_AFTER_STARTUP,
52 : : .request_fn = test_shmem_request,
53 : : .init_fn = test_shmem_init,
54 : : .attach_fn = test_shmem_attach,
55 : : };
56 : :
57 : : static void
58 : 16 : test_shmem_request(void *arg)
59 : : {
60 [ + - ]: 16 : elog(LOG, "test_shmem_request callback called");
61 : :
62 : 16 : ShmemRequestStruct(.name = "test_shmem area",
63 : : .size = test_shmem_area_size,
64 : : .ptr = (void **) &TestShmem);
65 : 15 : }
66 : :
67 : : static void
68 : 7 : test_shmem_init(void *arg)
69 : : {
70 [ + - ]: 7 : elog(LOG, "init callback called");
71 : :
45 72 : 7 : INJECTION_POINT("test-shmem-init", NULL);
73 : :
189 74 [ - + ]: 6 : if (TestShmem->initialized)
189 heikki.linnakangas@i 75 [ # # ]:UBC 0 : elog(ERROR, "shmem area already initialized");
189 heikki.linnakangas@i 76 :CBC 6 : TestShmem->initialized = true;
77 : :
78 [ - + ]: 6 : if (attached_or_initialized)
189 heikki.linnakangas@i 79 [ # # ]:UBC 0 : elog(ERROR, "attach or initialize already called in this process");
189 heikki.linnakangas@i 80 :CBC 6 : attached_or_initialized = true;
81 : 6 : }
82 : :
83 : : static void
84 : 6 : test_shmem_attach(void *arg)
85 : : {
86 [ + - ]: 6 : elog(LOG, "test_shmem_attach callback called");
87 [ - + ]: 6 : if (!TestShmem->initialized)
189 heikki.linnakangas@i 88 [ # # ]:UBC 0 : elog(ERROR, "shmem area not yet initialized");
189 heikki.linnakangas@i 89 :CBC 6 : TestShmem->attach_count++;
90 : :
91 [ - + ]: 6 : if (attached_or_initialized)
189 heikki.linnakangas@i 92 [ # # ]:UBC 0 : elog(ERROR, "attach or initialize already called in this process");
189 heikki.linnakangas@i 93 :CBC 6 : attached_or_initialized = true;
94 : 6 : }
95 : :
96 : : void
97 : 16 : _PG_init(void)
98 : : {
99 [ + - ]: 16 : elog(LOG, "test_shmem module's _PG_init called");
100 : :
45 101 [ + + ]: 16 : if (!test_shmem_guc_defined)
102 : : {
103 : : /*
104 : : * The minimum size that makes sense is sizeof(TestShmemData), but we
105 : : * allow -1 so that we can test passing SHMEM_ATTACH_UNKNOWN_SIZE.
106 : : */
107 : 15 : DefineCustomIntVariable("test_shmem.area_size",
108 : : "Size of the shmem area to request.",
109 : : NULL,
110 : : &test_shmem_area_size,
111 : : DEFAULT_TEST_AREA_BYTES,
112 : : -1,
113 : : MAX_TEST_AREA_BYTES,
114 : : PGC_USERSET,
115 : : GUC_UNIT_BYTE,
116 : : NULL, NULL, NULL);
117 : 15 : MarkGUCPrefixReserved("test_shmem");
118 : 15 : test_shmem_guc_defined = true;
119 : : }
189 120 : 16 : RegisterShmemCallbacks(&TestShmemCallbacks);
121 : 13 : }
122 : :
123 : 7 : PG_FUNCTION_INFO_V1(get_test_shmem_attach_count);
124 : : Datum
125 : 6 : get_test_shmem_attach_count(PG_FUNCTION_ARGS)
126 : : {
127 [ - + ]: 6 : if (!attached_or_initialized)
189 heikki.linnakangas@i 128 [ # # ]:UBC 0 : elog(ERROR, "shmem area not attached or initialized in this process");
189 heikki.linnakangas@i 129 [ - + ]:CBC 6 : if (!TestShmem->initialized)
189 heikki.linnakangas@i 130 [ # # ]:UBC 0 : elog(ERROR, "shmem area not yet initialized");
189 heikki.linnakangas@i 131 :CBC 6 : PG_RETURN_INT32(TestShmem->attach_count);
132 : : }
133 : :
134 : :
135 : : /*
136 : : * Callback for test_shmem_register(). test_shmem_register() provides the
137 : : * options, we just pass them through to ShmemRequestStruct.
138 : : */
139 : : static void
12 140 : 8 : test_shmem_after_startup_request(void *arg)
141 : : {
142 : 8 : ShmemStructOpts *opts = (ShmemStructOpts *) arg;
143 : :
144 [ + - ]: 8 : elog(LOG, "test_shmem_after_startup_request callback called");
145 : :
146 : 8 : ShmemRequestStructWithOpts(opts);
147 : 8 : }
148 : :
149 : : /*
150 : : * Allocate or attach to a shmem structure, with the caller-supplied name and
151 : : * size.
152 : : *
153 : : * The given integer 'new_value' is stored in the area, and the old value
154 : : * is returned.
155 : : */
156 : 7 : PG_FUNCTION_INFO_V1(test_shmem_register);
157 : : Datum
158 : 8 : test_shmem_register(PG_FUNCTION_ARGS)
159 : : {
160 : 8 : char *name = text_to_cstring(PG_GETARG_TEXT_PP(0));
161 : 8 : int64 size = PG_GETARG_INT64(1);
162 : 8 : int new_value = PG_GETARG_INT32(2);
163 : : int old_value;
164 : 8 : int *attached = NULL;
165 : :
166 : 8 : ShmemStructOpts opts = {
167 : : .name = name,
168 : : .size = size,
169 : : .ptr = (void **) &attached,
170 : : };
171 : :
172 : 8 : ShmemCallbacks callbacks = {
173 : : .flags = SHMEM_CALLBACKS_ALLOW_AFTER_STARTUP,
174 : : .request_fn = test_shmem_after_startup_request,
175 : : .opaque_arg = &opts,
176 : : };
177 : :
178 : 8 : RegisterShmemCallbacks(&callbacks);
179 [ - + ]: 6 : if (attached == NULL)
12 heikki.linnakangas@i 180 [ # # ]:UBC 0 : elog(ERROR, "could not attach to shared memory");
181 : :
12 heikki.linnakangas@i 182 :CBC 6 : old_value = *attached;
183 : 6 : *attached = new_value;
184 : :
185 : 6 : PG_RETURN_INT32(old_value);
186 : : }
|