Branch data Line data Source code
1 : : /*-------------------------------------------------------------------------
2 : : *
3 : : * stashfuncs.c
4 : : * SQL interface to pg_stash_advice
5 : : *
6 : : * Copyright (c) 2016-2026, PostgreSQL Global Development Group
7 : : *
8 : : * contrib/pg_stash_advice/stashfuncs.c
9 : : *
10 : : *-------------------------------------------------------------------------
11 : : */
12 : : #include "postgres.h"
13 : :
14 : : #include "common/hashfn.h"
15 : : #include "fmgr.h"
16 : : #include "funcapi.h"
17 : : #include "miscadmin.h"
18 : : #include "pg_stash_advice.h"
19 : : #include "utils/builtins.h"
20 : : #include "utils/tuplestore.h"
21 : :
22 : 6 : PG_FUNCTION_INFO_V1(pg_create_advice_stash);
23 : 4 : PG_FUNCTION_INFO_V1(pg_drop_advice_stash);
24 : 4 : PG_FUNCTION_INFO_V1(pg_get_advice_stash_contents);
25 : 7 : PG_FUNCTION_INFO_V1(pg_get_advice_stashes);
26 : 4 : PG_FUNCTION_INFO_V1(pg_set_stashed_advice);
27 : 2 : PG_FUNCTION_INFO_V1(pg_start_stash_advice_worker);
28 : :
29 : : typedef struct pgsa_stash_count
30 : : {
31 : : uint32 status;
32 : : uint64 pgsa_stash_id;
33 : : int64 num_entries;
34 : : } pgsa_stash_count;
35 : :
36 : : #define SH_PREFIX pgsa_stash_count_table
37 : : #define SH_ELEMENT_TYPE pgsa_stash_count
38 : : #define SH_KEY_TYPE uint64
39 : : #define SH_KEY pgsa_stash_id
40 : : #define SH_HASH_KEY(tb, key) hash_bytes((const unsigned char *) &(key), sizeof(uint64))
41 : : #define SH_EQUAL(tb, a, b) (a == b)
42 : : #define SH_SCOPE static inline
43 : : #define SH_DEFINE
44 : : #define SH_DECLARE
45 : : #include "lib/simplehash.h"
46 : :
47 : : /*
48 : : * SQL-callable function to create an advice stash
49 : : */
50 : : Datum
51 : 10 : pg_create_advice_stash(PG_FUNCTION_ARGS)
52 : : {
53 : 10 : char *stash_name = text_to_cstring(PG_GETARG_TEXT_PP(0));
54 : :
55 : 10 : pgsa_check_stash_name(stash_name);
56 [ + + ]: 6 : if (unlikely(pgsa_entry_dshash == NULL))
57 : 3 : pgsa_attach();
58 : 6 : pgsa_check_lockout();
59 : 6 : LWLockAcquire(&pgsa_state->lock, LW_EXCLUSIVE);
60 : 6 : pgsa_create_stash(stash_name);
61 : 5 : LWLockRelease(&pgsa_state->lock);
62 : 5 : PG_RETURN_VOID();
63 : : }
64 : :
65 : : /*
66 : : * SQL-callable function to drop an advice stash
67 : : */
68 : : Datum
69 : 6 : pg_drop_advice_stash(PG_FUNCTION_ARGS)
70 : : {
71 : 6 : char *stash_name = text_to_cstring(PG_GETARG_TEXT_PP(0));
72 : :
73 : 6 : pgsa_check_stash_name(stash_name);
74 [ + + ]: 6 : if (unlikely(pgsa_entry_dshash == NULL))
75 : 1 : pgsa_attach();
76 : 6 : pgsa_check_lockout();
77 : 6 : LWLockAcquire(&pgsa_state->lock, LW_EXCLUSIVE);
78 : 6 : pgsa_drop_stash(stash_name);
79 : 5 : LWLockRelease(&pgsa_state->lock);
80 : 5 : PG_RETURN_VOID();
81 : : }
82 : :
83 : : /*
84 : : * SQL-callable function to provide a list of advice stashes
85 : : */
86 : : Datum
87 : 6 : pg_get_advice_stashes(PG_FUNCTION_ARGS)
88 : : {
89 : 6 : ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
90 : : dshash_seq_status iterator;
91 : : pgsa_entry *entry;
92 : : pgsa_stash *stash;
93 : : pgsa_stash_count_table_hash *chash;
94 : :
95 : 6 : InitMaterializedSRF(fcinfo, 0);
96 : :
97 : : /* Attach to dynamic shared memory if not already done. */
98 [ + + ]: 6 : if (unlikely(pgsa_entry_dshash == NULL))
99 : 4 : pgsa_attach();
100 : :
101 : : /* If stash data is still being restored from disk, ignore. */
102 [ - + ]: 6 : if (pg_atomic_unlocked_test_flag(&pgsa_state->stashes_ready))
103 : 0 : return (Datum) 0;
104 : :
105 : : /* Tally up the number of entries per stash. */
106 : 6 : chash = pgsa_stash_count_table_create(CurrentMemoryContext, 64, NULL);
107 : 6 : dshash_seq_init(&iterator, pgsa_entry_dshash, true);
108 [ + + ]: 18 : while ((entry = dshash_seq_next(&iterator)) != NULL)
109 : : {
110 : : pgsa_stash_count *c;
111 : : bool found;
112 : :
113 : 12 : c = pgsa_stash_count_table_insert(chash,
114 : : entry->key.pgsa_stash_id,
115 : : &found);
116 [ + + ]: 12 : if (!found)
117 : 8 : c->num_entries = 1;
118 : : else
119 : 4 : c->num_entries++;
120 : : }
121 : 6 : dshash_seq_term(&iterator);
122 : :
123 : : /* Emit results. */
124 : 6 : dshash_seq_init(&iterator, pgsa_stash_dshash, true);
125 [ + + ]: 17 : while ((stash = dshash_seq_next(&iterator)) != NULL)
126 : : {
127 : : Datum values[2];
128 : : bool nulls[2];
129 : : pgsa_stash_count *c;
130 : :
131 : 11 : values[0] = CStringGetTextDatum(stash->name);
132 : 11 : nulls[0] = false;
133 : :
134 : 11 : c = pgsa_stash_count_table_lookup(chash, stash->pgsa_stash_id);
135 [ + + ]: 11 : values[1] = Int64GetDatum(c == NULL ? 0 : c->num_entries);
136 : 11 : nulls[1] = false;
137 : :
138 : 11 : tuplestore_putvalues(rsinfo->setResult, rsinfo->setDesc, values,
139 : : nulls);
140 : : }
141 : 6 : dshash_seq_term(&iterator);
142 : :
143 : 6 : return (Datum) 0;
144 : : }
145 : :
146 : : /*
147 : : * SQL-callable function to provide advice stash contents
148 : : */
149 : : Datum
150 : 6 : pg_get_advice_stash_contents(PG_FUNCTION_ARGS)
151 : : {
152 : 6 : ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
153 : : dshash_seq_status iterator;
154 : 6 : char *stash_name = NULL;
155 : 6 : pgsa_stash_name_table_hash *nhash = NULL;
156 : 6 : uint64 stash_id = 0;
157 : : pgsa_entry *entry;
158 : :
159 : 6 : InitMaterializedSRF(fcinfo, 0);
160 : :
161 : : /* Attach to dynamic shared memory if not already done. */
162 [ + + ]: 6 : if (unlikely(pgsa_entry_dshash == NULL))
163 : 1 : pgsa_attach();
164 : :
165 : : /* If stash data is still being restored from disk, ignore. */
166 [ - + ]: 6 : if (pg_atomic_unlocked_test_flag(&pgsa_state->stashes_ready))
167 : 0 : return (Datum) 0;
168 : :
169 : : /* User can pass NULL for all stashes, or the name of a specific stash. */
170 [ + + ]: 6 : if (!PG_ARGISNULL(0))
171 : : {
172 : 4 : stash_name = text_to_cstring(PG_GETARG_TEXT_PP(0));
173 : 4 : pgsa_check_stash_name(stash_name);
174 : 4 : stash_id = pgsa_lookup_stash_id(stash_name);
175 : :
176 : : /* If the user specified a stash name, it should exist. */
177 [ + + ]: 4 : if (stash_id == 0)
178 [ + - ]: 1 : ereport(ERROR,
179 : : errcode(ERRCODE_INVALID_PARAMETER_VALUE),
180 : : errmsg("advice stash \"%s\" does not exist", stash_name));
181 : : }
182 : : else
183 : : {
184 : : pgsa_stash *stash;
185 : :
186 : : /*
187 : : * If we're dumping data about all stashes, we need an ID->name lookup
188 : : * table.
189 : : */
190 : 2 : nhash = pgsa_stash_name_table_create(CurrentMemoryContext, 64, NULL);
191 : 2 : dshash_seq_init(&iterator, pgsa_stash_dshash, true);
192 [ + + ]: 6 : while ((stash = dshash_seq_next(&iterator)) != NULL)
193 : : {
194 : : pgsa_stash_name *n;
195 : : bool found;
196 : :
197 : 4 : n = pgsa_stash_name_table_insert(nhash,
198 : : stash->pgsa_stash_id,
199 : : &found);
200 : : Assert(!found);
201 : 4 : n->name = pstrdup(stash->name);
202 : : }
203 : 2 : dshash_seq_term(&iterator);
204 : : }
205 : :
206 : : /* Now iterate over all the entries. */
207 : 5 : dshash_seq_init(&iterator, pgsa_entry_dshash, false);
208 [ + + ]: 15 : while ((entry = dshash_seq_next(&iterator)) != NULL)
209 : : {
210 : : Datum values[3];
211 : : bool nulls[3];
212 : : char *this_stash_name;
213 : : char *advice_string;
214 : :
215 : : /* Skip incomplete entries where the advice string was never set. */
216 [ - + ]: 10 : if (entry->advice_string == InvalidDsaPointer)
217 : 2 : continue;
218 : :
219 [ + + ]: 10 : if (stash_id != 0)
220 : : {
221 : : /*
222 : : * We're only dumping data for one particular stash, so skip
223 : : * entries for any other stash and use the stash name specified by
224 : : * the user.
225 : : */
226 [ + + ]: 5 : if (stash_id != entry->key.pgsa_stash_id)
227 : 2 : continue;
228 : 3 : this_stash_name = stash_name;
229 : : }
230 : : else
231 : : {
232 : : pgsa_stash_name *n;
233 : :
234 : : /*
235 : : * We're dumping data for all stashes, so look up the correct name
236 : : * to use in the hash table. If nothing is found, which is
237 : : * possible due to race conditions, make up a string to use.
238 : : */
239 : 5 : n = pgsa_stash_name_table_lookup(nhash, entry->key.pgsa_stash_id);
240 [ + - ]: 5 : if (n != NULL)
241 : 5 : this_stash_name = n->name;
242 : : else
243 : 0 : this_stash_name = psprintf("<stash %" PRIu64 ">",
244 : : entry->key.pgsa_stash_id);
245 : : }
246 : :
247 : : /* Work out tuple values. */
248 : 8 : values[0] = CStringGetTextDatum(this_stash_name);
249 : 8 : nulls[0] = false;
250 : 8 : values[1] = Int64GetDatum(entry->key.queryId);
251 : 8 : nulls[1] = false;
252 : 8 : advice_string = dsa_get_address(pgsa_dsa_area, entry->advice_string);
253 : 8 : values[2] = CStringGetTextDatum(advice_string);
254 : 8 : nulls[2] = false;
255 : :
256 : : /* Emit the tuple. */
257 : 8 : tuplestore_putvalues(rsinfo->setResult, rsinfo->setDesc, values,
258 : : nulls);
259 : : }
260 : 5 : dshash_seq_term(&iterator);
261 : :
262 : 5 : return (Datum) 0;
263 : : }
264 : :
265 : : /*
266 : : * SQL-callable function to update an advice stash entry for a particular
267 : : * query ID
268 : : *
269 : : * If the advice string (the third argument) is NULL, we delete any existing
270 : : * advice stash entry; otherwise, we either create an entry or update it with
271 : : * the new advice string.
272 : : */
273 : : Datum
274 : 11 : pg_set_stashed_advice(PG_FUNCTION_ARGS)
275 : : {
276 : : char *stash_name;
277 : : int64 queryId;
278 : :
279 [ + - - + ]: 11 : if (PG_ARGISNULL(0) || PG_ARGISNULL(1))
280 : 0 : PG_RETURN_NULL();
281 : :
282 : : /* Get and check advice stash name. */
283 : 11 : stash_name = text_to_cstring(PG_GETARG_TEXT_PP(0));
284 : 11 : pgsa_check_stash_name(stash_name);
285 : :
286 : : /*
287 : : * Get and check query ID.
288 : : *
289 : : * Query ID 0 means no query ID was computed, so reject that.
290 : : */
291 : 11 : queryId = PG_GETARG_INT64(1);
292 [ + + ]: 11 : if (queryId == 0)
293 [ + - ]: 1 : ereport(ERROR,
294 : : errcode(ERRCODE_INVALID_PARAMETER_VALUE),
295 : : errmsg("cannot set advice string for query ID 0"));
296 : :
297 : : /* Attach to dynamic shared memory if not already done. */
298 [ - + ]: 10 : if (unlikely(pgsa_entry_dshash == NULL))
299 : 0 : pgsa_attach();
300 : :
301 : : /* Don't allow writes if stash data is still being restored from disk. */
302 : 10 : pgsa_check_lockout();
303 : :
304 : : /* Now call the appropriate function to do the real work. */
305 [ + + ]: 10 : if (PG_ARGISNULL(2))
306 : : {
307 : 2 : LWLockAcquire(&pgsa_state->lock, LW_SHARED);
308 : 2 : pgsa_clear_advice_string(stash_name, queryId);
309 : 1 : LWLockRelease(&pgsa_state->lock);
310 : : }
311 : : else
312 : : {
313 : 8 : char *advice_string = text_to_cstring(PG_GETARG_TEXT_PP(2));
314 : :
315 : 8 : LWLockAcquire(&pgsa_state->lock, LW_SHARED);
316 : 8 : pgsa_set_advice_string(stash_name, queryId, advice_string);
317 : 7 : LWLockRelease(&pgsa_state->lock);
318 : : }
319 : :
320 : 8 : PG_RETURN_VOID();
321 : : }
322 : :
323 : : /*
324 : : * SQL-callable function to start the persistence background worker.
325 : : */
326 : : Datum
327 : 0 : pg_start_stash_advice_worker(PG_FUNCTION_ARGS)
328 : : {
329 : : pid_t pid;
330 : :
331 [ # # ]: 0 : if (unlikely(pgsa_entry_dshash == NULL))
332 : 0 : pgsa_attach();
333 : :
334 : 0 : LWLockAcquire(&pgsa_state->lock, LW_SHARED);
335 : 0 : pid = pgsa_state->bgworker_pid;
336 : 0 : LWLockRelease(&pgsa_state->lock);
337 : :
338 [ # # ]: 0 : if (pid != InvalidPid)
339 [ # # ]: 0 : ereport(ERROR,
340 : : (errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
341 : : errmsg("pg_stash_advice worker is already running under PID %d",
342 : : (int) pid)));
343 : :
344 : 0 : pgsa_start_worker();
345 : :
346 : 0 : PG_RETURN_VOID();
347 : : }
|
