Branch data Line data Source code
1 : : /*-------------------------------------------------------------------------
2 : : *
3 : : * pg_stash_advice.c
4 : : * core infrastructure for pg_stash_advice contrib module
5 : : *
6 : : * Copyright (c) 2016-2026, PostgreSQL Global Development Group
7 : : *
8 : : * contrib/pg_stash_advice/pg_stash_advice.c
9 : : *
10 : : *-------------------------------------------------------------------------
11 : : */
12 : : #include "postgres.h"
13 : :
14 : : #include "common/hashfn.h"
15 : : #include "common/string.h"
16 : : #include "miscadmin.h"
17 : : #include "nodes/queryjumble.h"
18 : : #include "pg_plan_advice.h"
19 : : #include "pg_stash_advice.h"
20 : : #include "postmaster/bgworker.h"
21 : : #include "storage/dsm_registry.h"
22 : : #include "utils/guc.h"
23 : : #include "utils/memutils.h"
24 : :
25 : 6 : PG_MODULE_MAGIC_EXT(
26 : : .name = "pg_stash_advice",
27 : : .version = PG_VERSION
28 : : );
29 : :
30 : : /* Shared memory hash table parameters */
31 : : static dshash_parameters pgsa_stash_dshash_parameters = {
32 : : NAMEDATALEN,
33 : : sizeof(pgsa_stash),
34 : : dshash_strcmp,
35 : : dshash_strhash,
36 : : dshash_strcpy,
37 : : LWTRANCHE_INVALID /* gets set at runtime */
38 : : };
39 : :
40 : : static dshash_parameters pgsa_entry_dshash_parameters = {
41 : : sizeof(pgsa_entry_key),
42 : : sizeof(pgsa_entry),
43 : : dshash_memcmp,
44 : : dshash_memhash,
45 : : dshash_memcpy,
46 : : LWTRANCHE_INVALID /* gets set at runtime */
47 : : };
48 : :
49 : : /* GUC variables */
50 : : static char *pg_stash_advice_stash_name = "";
51 : : bool pg_stash_advice_persist = true;
52 : : int pg_stash_advice_persist_interval = 30;
53 : :
54 : : /* Shared memory pointers */
55 : : pgsa_shared_state *pgsa_state;
56 : : dsa_area *pgsa_dsa_area;
57 : : dshash_table *pgsa_stash_dshash;
58 : : dshash_table *pgsa_entry_dshash;
59 : :
60 : : /* Other global variables */
61 : : static MemoryContext pg_stash_advice_mcxt;
62 : :
63 : : /* Function prototypes */
64 : : static char *pgsa_advisor(PlannerGlobal *glob,
65 : : Query *parse,
66 : : const char *query_string,
67 : : int cursorOptions,
68 : : ExplainState *es);
69 : : static bool pgsa_check_stash_name_guc(char **newval, void **extra,
70 : : GucSource source);
71 : : static void pgsa_init_shared_state(void *ptr, void *arg);
72 : : static bool pgsa_is_identifier(char *str);
73 : :
74 : : /* Stash name -> stash ID hash table */
75 : : #define SH_PREFIX pgsa_stash_name_table
76 : : #define SH_ELEMENT_TYPE pgsa_stash_name
77 : : #define SH_KEY_TYPE uint64
78 : : #define SH_KEY pgsa_stash_id
79 : : #define SH_HASH_KEY(tb, key) hash_bytes((const unsigned char *) &(key), sizeof(uint64))
80 : : #define SH_EQUAL(tb, a, b) (a == b)
81 : : #define SH_SCOPE extern
82 : : #define SH_DEFINE
83 : : #include "lib/simplehash.h"
84 : :
85 : : /*
86 : : * Initialize this module.
87 : : */
88 : : void
89 : 6 : _PG_init(void)
90 : : {
91 : : void (*add_advisor_fn) (pg_plan_advice_advisor_hook hook);
92 : :
93 : : /* If compute_query_id = 'auto', we would like query IDs. */
94 : 6 : EnableQueryId();
95 : :
96 : : /* Define our GUCs. */
97 [ + + ]: 6 : if (process_shared_preload_libraries_in_progress)
98 : 4 : DefineCustomBoolVariable("pg_stash_advice.persist",
99 : : "Save and restore advice stash contents across restarts.",
100 : : NULL,
101 : : &pg_stash_advice_persist,
102 : : true,
103 : : PGC_POSTMASTER,
104 : : 0,
105 : : NULL,
106 : : NULL,
107 : : NULL);
108 : : else
109 : 2 : pg_stash_advice_persist = false;
110 : :
111 : 6 : DefineCustomIntVariable("pg_stash_advice.persist_interval",
112 : : "Interval between advice stash saves, in seconds.",
113 : : NULL,
114 : : &pg_stash_advice_persist_interval,
115 : : 30,
116 : : 0,
117 : : 3600,
118 : : PGC_SIGHUP,
119 : : GUC_UNIT_S,
120 : : NULL,
121 : : NULL,
122 : : NULL);
123 : :
124 : 6 : DefineCustomStringVariable("pg_stash_advice.stash_name",
125 : : "Name of the advice stash to be used in this session.",
126 : : NULL,
127 : : &pg_stash_advice_stash_name,
128 : : "",
129 : : PGC_USERSET,
130 : : 0,
131 : : pgsa_check_stash_name_guc,
132 : : NULL,
133 : : NULL);
134 : :
135 : 6 : MarkGUCPrefixReserved("pg_stash_advice");
136 : :
137 : : /* Start the background worker for persistence, if enabled. */
138 [ + + ]: 6 : if (pg_stash_advice_persist)
139 : 4 : pgsa_start_worker();
140 : :
141 : : /* Tell pg_plan_advice that we want to provide advice strings. */
142 : 6 : add_advisor_fn =
143 : 6 : load_external_function("pg_plan_advice", "pg_plan_advice_add_advisor",
144 : : true, NULL);
145 : 6 : (*add_advisor_fn) (pgsa_advisor);
146 : 6 : }
147 : :
148 : : /*
149 : : * Get the advice string that has been configured for this query, if any,
150 : : * and return it. Otherwise, return NULL.
151 : : */
152 : : static char *
153 : 53 : pgsa_advisor(PlannerGlobal *glob, Query *parse,
154 : : const char *query_string, int cursorOptions,
155 : : ExplainState *es)
156 : : {
157 : : pgsa_entry_key key;
158 : : pgsa_entry *entry;
159 : : char *advice_string;
160 : : uint64 stash_id;
161 : :
162 : : /*
163 : : * Exit quickly if the stash name is empty or there's no query ID.
164 : : */
165 [ + + - + ]: 53 : if (pg_stash_advice_stash_name[0] == '\0' || parse->queryId == 0)
166 : 23 : return NULL;
167 : :
168 : : /* Attach to dynamic shared memory if not already done. */
169 [ - + ]: 30 : if (unlikely(pgsa_entry_dshash == NULL))
170 : 0 : pgsa_attach();
171 : :
172 : : /* If stash data is still being restored from disk, ignore. */
173 [ - + ]: 30 : if (pg_atomic_unlocked_test_flag(&pgsa_state->stashes_ready))
174 : 0 : return NULL;
175 : :
176 : : /*
177 : : * Translate pg_stash_advice.stash_name to an integer ID.
178 : : *
179 : : * pgsa_check_stash_name_guc() has already validated the advice stash
180 : : * name, so we don't need to call pgsa_check_stash_name() here.
181 : : */
182 : 30 : stash_id = pgsa_lookup_stash_id(pg_stash_advice_stash_name);
183 [ + + ]: 30 : if (stash_id == 0)
184 : 1 : return NULL;
185 : :
186 : : /*
187 : : * Look up the advice string for the given stash ID + query ID.
188 : : *
189 : : * If we find an advice string, we copy it into the current memory
190 : : * context, presumably short-lived, so that we can release the lock on the
191 : : * dshash entry. pg_plan_advice only needs the value to remain allocated
192 : : * long enough for it to be parsed, so this should be good enough.
193 : : */
194 : 29 : memset(&key, 0, sizeof(pgsa_entry_key));
195 : 29 : key.pgsa_stash_id = stash_id;
196 : 29 : key.queryId = parse->queryId;
197 : 29 : entry = dshash_find(pgsa_entry_dshash, &key, false);
198 [ + + ]: 29 : if (entry == NULL)
199 : 25 : return NULL;
200 [ - + ]: 4 : if (entry->advice_string == InvalidDsaPointer)
201 : 0 : advice_string = NULL;
202 : : else
203 : 4 : advice_string = pstrdup(dsa_get_address(pgsa_dsa_area,
204 : : entry->advice_string));
205 : 4 : dshash_release_lock(pgsa_entry_dshash, entry);
206 : :
207 : : /* If we found an advice string, emit a debug message. */
208 [ + - ]: 4 : if (advice_string != NULL)
209 [ - + ]: 4 : elog(DEBUG2, "supplying automatic advice for stash \"%s\", query ID %" PRId64 ": %s",
210 : : pg_stash_advice_stash_name, key.queryId, advice_string);
211 : :
212 : 4 : return advice_string;
213 : : }
214 : :
215 : : /*
216 : : * Attach to various structures in dynamic shared memory.
217 : : *
218 : : * This function is designed to be resilient against errors. That is, if it
219 : : * fails partway through, it should be possible to call it again, repeat no
220 : : * work already completed, and potentially succeed or at least get further if
221 : : * whatever caused the previous failure has been corrected.
222 : : */
223 : : void
224 : 13 : pgsa_attach(void)
225 : : {
226 : : bool found;
227 : : MemoryContext oldcontext;
228 : :
229 : : /*
230 : : * Create a memory context to make sure that any control structures
231 : : * allocated in local memory are sufficiently persistent.
232 : : */
233 [ + - ]: 13 : if (pg_stash_advice_mcxt == NULL)
234 : 13 : pg_stash_advice_mcxt = AllocSetContextCreate(TopMemoryContext,
235 : : "pg_stash_advice",
236 : : ALLOCSET_DEFAULT_SIZES);
237 : 13 : oldcontext = MemoryContextSwitchTo(pg_stash_advice_mcxt);
238 : :
239 : : /* Attach to the fixed-size state object if not already done. */
240 [ + - ]: 13 : if (pgsa_state == NULL)
241 : 13 : pgsa_state = GetNamedDSMSegment("pg_stash_advice",
242 : : sizeof(pgsa_shared_state),
243 : : pgsa_init_shared_state,
244 : : &found, NULL);
245 : :
246 : : /* Attach to the DSA area if not already done. */
247 [ + - ]: 13 : if (pgsa_dsa_area == NULL)
248 : : {
249 : : dsa_handle area_handle;
250 : :
251 : 13 : LWLockAcquire(&pgsa_state->lock, LW_EXCLUSIVE);
252 : 13 : area_handle = pgsa_state->area;
253 [ + + ]: 13 : if (area_handle == DSA_HANDLE_INVALID)
254 : : {
255 : 5 : pgsa_dsa_area = dsa_create(pgsa_state->dsa_tranche);
256 : 5 : dsa_pin(pgsa_dsa_area);
257 : 5 : pgsa_state->area = dsa_get_handle(pgsa_dsa_area);
258 : 5 : LWLockRelease(&pgsa_state->lock);
259 : : }
260 : : else
261 : : {
262 : 8 : LWLockRelease(&pgsa_state->lock);
263 : 8 : pgsa_dsa_area = dsa_attach(area_handle);
264 : : }
265 : 13 : dsa_pin_mapping(pgsa_dsa_area);
266 : : }
267 : :
268 : : /* Attach to the stash_name->stash_id hash table if not already done. */
269 [ + - ]: 13 : if (pgsa_stash_dshash == NULL)
270 : : {
271 : : dshash_table_handle stash_handle;
272 : :
273 : 13 : LWLockAcquire(&pgsa_state->lock, LW_EXCLUSIVE);
274 : 13 : pgsa_stash_dshash_parameters.tranche_id = pgsa_state->stash_tranche;
275 : 13 : stash_handle = pgsa_state->stash_hash;
276 [ + + ]: 13 : if (stash_handle == DSHASH_HANDLE_INVALID)
277 : : {
278 : 5 : pgsa_stash_dshash = dshash_create(pgsa_dsa_area,
279 : : &pgsa_stash_dshash_parameters,
280 : : NULL);
281 : 10 : pgsa_state->stash_hash =
282 : 5 : dshash_get_hash_table_handle(pgsa_stash_dshash);
283 : 5 : LWLockRelease(&pgsa_state->lock);
284 : : }
285 : : else
286 : : {
287 : 8 : LWLockRelease(&pgsa_state->lock);
288 : 8 : pgsa_stash_dshash = dshash_attach(pgsa_dsa_area,
289 : : &pgsa_stash_dshash_parameters,
290 : : stash_handle, NULL);
291 : : }
292 : : }
293 : :
294 : : /* Attach to the entry hash table if not already done. */
295 [ + - ]: 13 : if (pgsa_entry_dshash == NULL)
296 : : {
297 : : dshash_table_handle entry_handle;
298 : :
299 : 13 : LWLockAcquire(&pgsa_state->lock, LW_EXCLUSIVE);
300 : 13 : pgsa_entry_dshash_parameters.tranche_id = pgsa_state->entry_tranche;
301 : 13 : entry_handle = pgsa_state->entry_hash;
302 [ + + ]: 13 : if (entry_handle == DSHASH_HANDLE_INVALID)
303 : : {
304 : 5 : pgsa_entry_dshash = dshash_create(pgsa_dsa_area,
305 : : &pgsa_entry_dshash_parameters,
306 : : NULL);
307 : 10 : pgsa_state->entry_hash =
308 : 5 : dshash_get_hash_table_handle(pgsa_entry_dshash);
309 : 5 : LWLockRelease(&pgsa_state->lock);
310 : : }
311 : : else
312 : : {
313 : 8 : LWLockRelease(&pgsa_state->lock);
314 : 8 : pgsa_entry_dshash = dshash_attach(pgsa_dsa_area,
315 : : &pgsa_entry_dshash_parameters,
316 : : entry_handle, NULL);
317 : : }
318 : : }
319 : :
320 : : /* Restore previous memory context. */
321 : 13 : MemoryContextSwitchTo(oldcontext);
322 : 13 : }
323 : :
324 : : /*
325 : : * Error out if the stashes have not been loaded from disk yet.
326 : : */
327 : : void
328 : 22 : pgsa_check_lockout(void)
329 : : {
330 [ - + ]: 22 : if (pg_atomic_unlocked_test_flag(&pgsa_state->stashes_ready))
331 [ # # ]: 0 : ereport(ERROR,
332 : : (errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
333 : : errmsg("stash modifications are not allowed because \"%s\" has not been loaded yet",
334 : : PGSA_DUMP_FILE)));
335 : 22 : }
336 : :
337 : : /*
338 : : * Check whether an advice stash name is legal, and signal an error if not.
339 : : *
340 : : * Keep this in sync with pgsa_check_stash_name_guc, below.
341 : : */
342 : : void
343 : 31 : pgsa_check_stash_name(char *stash_name)
344 : : {
345 : : /* Reject empty advice stash name. */
346 [ + + ]: 31 : if (stash_name[0] == '\0')
347 [ + - ]: 1 : ereport(ERROR,
348 : : errcode(ERRCODE_INVALID_PARAMETER_VALUE),
349 : : errmsg("advice stash name may not be zero length"));
350 : :
351 : : /* Reject overlong advice stash names. */
352 [ + + ]: 30 : if (strlen(stash_name) + 1 > NAMEDATALEN)
353 [ + - ]: 1 : ereport(ERROR,
354 : : errcode(ERRCODE_INVALID_PARAMETER_VALUE),
355 : : errmsg("advice stash names may not be longer than %d bytes",
356 : : NAMEDATALEN - 1));
357 : :
358 : : /*
359 : : * Reject non-ASCII advice stash names, since advice stashes are visible
360 : : * across all databases and the encodings of those databases might differ.
361 : : */
362 [ + + ]: 29 : if (!pg_is_ascii(stash_name))
363 [ + - ]: 1 : ereport(ERROR,
364 : : errcode(ERRCODE_INVALID_PARAMETER_VALUE),
365 : : errmsg("advice stash name must not contain non-ASCII characters"));
366 : :
367 : : /*
368 : : * Reject things that do not look like identifiers, since the ability to
369 : : * create an advice stash with non-printable characters or weird symbols
370 : : * in the name is not likely to be useful to anyone.
371 : : */
372 [ + + ]: 28 : if (!pgsa_is_identifier(stash_name))
373 [ + - ]: 1 : ereport(ERROR,
374 : : errcode(ERRCODE_INVALID_PARAMETER_VALUE),
375 : : errmsg("advice stash name must begin with a letter or underscore and contain only letters, digits, and underscores"));
376 : 27 : }
377 : :
378 : : /*
379 : : * As above, but for the GUC check_hook. We allow the empty string here,
380 : : * though, as equivalent to disabling the feature.
381 : : */
382 : : static bool
383 : 11 : pgsa_check_stash_name_guc(char **newval, void **extra, GucSource source)
384 : : {
385 : 11 : char *stash_name = *newval;
386 : :
387 : : /* Reject overlong advice stash names. */
388 [ - + ]: 11 : if (strlen(stash_name) + 1 > NAMEDATALEN)
389 : : {
390 : 0 : GUC_check_errcode(ERRCODE_INVALID_PARAMETER_VALUE);
391 : 0 : GUC_check_errdetail("Advice stash names may not be longer than %d bytes.",
392 : : NAMEDATALEN - 1);
393 : 0 : return false;
394 : : }
395 : :
396 : : /*
397 : : * Reject non-ASCII advice stash names, since advice stashes are visible
398 : : * across all databases and the encodings of those databases might differ.
399 : : */
400 [ + + ]: 11 : if (!pg_is_ascii(stash_name))
401 : : {
402 : 1 : GUC_check_errcode(ERRCODE_INVALID_PARAMETER_VALUE);
403 : 1 : GUC_check_errdetail("Advice stash name must not contain non-ASCII characters.");
404 : 1 : return false;
405 : : }
406 : :
407 : : /*
408 : : * Reject things that do not look like identifiers, since the ability to
409 : : * create an advice stash with non-printable characters or weird symbols
410 : : * in the name is not likely to be useful to anyone.
411 : : */
412 [ + + ]: 10 : if (!pgsa_is_identifier(stash_name))
413 : : {
414 : 1 : GUC_check_errcode(ERRCODE_INVALID_PARAMETER_VALUE);
415 : 1 : GUC_check_errdetail("Advice stash name must begin with a letter or underscore and contain only letters, digits, and underscores.");
416 : 1 : return false;
417 : : }
418 : :
419 : 9 : return true;
420 : : }
421 : :
422 : : /*
423 : : * Create an advice stash.
424 : : */
425 : : void
426 : 11 : pgsa_create_stash(char *stash_name)
427 : : {
428 : : pgsa_stash *stash;
429 : : bool found;
430 : :
431 : : Assert(LWLockHeldByMeInMode(&pgsa_state->lock, LW_EXCLUSIVE));
432 : :
433 : : /* Create a stash with this name, unless one already exists. */
434 : 11 : stash = dshash_find_or_insert(pgsa_stash_dshash, stash_name, &found);
435 [ + + ]: 11 : if (found)
436 [ + - ]: 1 : ereport(ERROR,
437 : : errcode(ERRCODE_INVALID_PARAMETER_VALUE),
438 : : errmsg("advice stash \"%s\" already exists", stash_name));
439 : 10 : stash->pgsa_stash_id = pgsa_state->next_stash_id++;
440 : 10 : dshash_release_lock(pgsa_stash_dshash, stash);
441 : :
442 : : /* Bump change count. */
443 : 10 : pg_atomic_add_fetch_u64(&pgsa_state->change_count, 1);
444 : 10 : }
445 : :
446 : : /*
447 : : * Remove any stored advice string for the given advice stash and query ID.
448 : : */
449 : : void
450 : 2 : pgsa_clear_advice_string(char *stash_name, int64 queryId)
451 : : {
452 : : pgsa_entry *entry;
453 : : pgsa_entry_key key;
454 : : uint64 stash_id;
455 : : dsa_pointer old_dp;
456 : :
457 : : Assert(LWLockHeldByMe(&pgsa_state->lock));
458 : :
459 : : /* Translate the stash name to an integer ID. */
460 [ + + ]: 2 : if ((stash_id = pgsa_lookup_stash_id(stash_name)) == 0)
461 [ + - ]: 1 : ereport(ERROR,
462 : : errcode(ERRCODE_INVALID_PARAMETER_VALUE),
463 : : errmsg("advice stash \"%s\" does not exist", stash_name));
464 : :
465 : : /*
466 : : * Look for an existing entry, and free it. But, be sure to save the
467 : : * pointer to the associated advice string, if any.
468 : : */
469 : 1 : memset(&key, 0, sizeof(pgsa_entry_key));
470 : 1 : key.pgsa_stash_id = stash_id;
471 : 1 : key.queryId = queryId;
472 : 1 : entry = dshash_find(pgsa_entry_dshash, &key, true);
473 [ - + ]: 1 : if (entry == NULL)
474 : 0 : old_dp = InvalidDsaPointer;
475 : : else
476 : : {
477 : 1 : old_dp = entry->advice_string;
478 : 1 : dshash_delete_entry(pgsa_entry_dshash, entry);
479 : : }
480 : :
481 : : /* Now we free the advice string as well, if there was one. */
482 [ + - ]: 1 : if (old_dp != InvalidDsaPointer)
483 : 1 : dsa_free(pgsa_dsa_area, old_dp);
484 : :
485 : : /* Bump change count. */
486 : 1 : pg_atomic_add_fetch_u64(&pgsa_state->change_count, 1);
487 : 1 : }
488 : :
489 : : /*
490 : : * Drop an advice stash.
491 : : */
492 : : void
493 : 6 : pgsa_drop_stash(char *stash_name)
494 : : {
495 : : pgsa_entry *entry;
496 : : pgsa_stash *stash;
497 : : dshash_seq_status iterator;
498 : : uint64 stash_id;
499 : :
500 : : Assert(LWLockHeldByMeInMode(&pgsa_state->lock, LW_EXCLUSIVE));
501 : :
502 : : /* Remove the entry for this advice stash. */
503 : 6 : stash = dshash_find(pgsa_stash_dshash, stash_name, true);
504 [ + + ]: 6 : if (stash == NULL)
505 [ + - ]: 1 : ereport(ERROR,
506 : : errcode(ERRCODE_INVALID_PARAMETER_VALUE),
507 : : errmsg("advice stash \"%s\" does not exist", stash_name));
508 : 5 : stash_id = stash->pgsa_stash_id;
509 : 5 : dshash_delete_entry(pgsa_stash_dshash, stash);
510 : :
511 : : /*
512 : : * Now remove all the entries. Since pgsa_state->lock must be held at
513 : : * least in shared mode to insert entries into pgsa_entry_dshash, it
514 : : * doesn't matter whether we do this before or after deleting the entry
515 : : * from pgsa_stash_dshash.
516 : : */
517 : 5 : dshash_seq_init(&iterator, pgsa_entry_dshash, true);
518 [ + + ]: 15 : while ((entry = dshash_seq_next(&iterator)) != NULL)
519 : : {
520 [ + + ]: 5 : if (stash_id == entry->key.pgsa_stash_id)
521 : : {
522 [ + - ]: 4 : if (entry->advice_string != InvalidDsaPointer)
523 : 4 : dsa_free(pgsa_dsa_area, entry->advice_string);
524 : 4 : dshash_delete_current(&iterator);
525 : : }
526 : : }
527 : 5 : dshash_seq_term(&iterator);
528 : :
529 : : /* Bump change count. */
530 : 5 : pg_atomic_add_fetch_u64(&pgsa_state->change_count, 1);
531 : 5 : }
532 : :
533 : : /*
534 : : * Remove all stashes and entries from shared memory.
535 : : *
536 : : * This is intended to be called before reloading from a dump file, so that
537 : : * a failed previous attempt doesn't leave stale data behind.
538 : : */
539 : : void
540 : 4 : pgsa_reset_all_stashes(void)
541 : : {
542 : : dshash_seq_status iter;
543 : : pgsa_entry *entry;
544 : :
545 : : Assert(LWLockHeldByMeInMode(&pgsa_state->lock, LW_EXCLUSIVE));
546 : :
547 : : /* Remove all stashes. */
548 : 4 : dshash_seq_init(&iter, pgsa_stash_dshash, true);
549 [ - + ]: 4 : while (dshash_seq_next(&iter) != NULL)
550 : 0 : dshash_delete_current(&iter);
551 : 4 : dshash_seq_term(&iter);
552 : :
553 : : /* Remove all entries. */
554 : 4 : dshash_seq_init(&iter, pgsa_entry_dshash, true);
555 [ - + ]: 4 : while ((entry = dshash_seq_next(&iter)) != NULL)
556 : : {
557 [ # # ]: 0 : if (entry->advice_string != InvalidDsaPointer)
558 : 0 : dsa_free(pgsa_dsa_area, entry->advice_string);
559 : 0 : dshash_delete_current(&iter);
560 : : }
561 : 4 : dshash_seq_term(&iter);
562 : :
563 : : /* Reset the stash ID counter. */
564 : 4 : pgsa_state->next_stash_id = UINT64CONST(1);
565 : 4 : }
566 : :
567 : : /*
568 : : * Initialize shared state when first created.
569 : : */
570 : : static void
571 : 5 : pgsa_init_shared_state(void *ptr, void *arg)
572 : : {
573 : 5 : pgsa_shared_state *state = (pgsa_shared_state *) ptr;
574 : :
575 : 5 : LWLockInitialize(&state->lock,
576 : : LWLockNewTrancheId("pg_stash_advice_lock"));
577 : 5 : state->dsa_tranche = LWLockNewTrancheId("pg_stash_advice_dsa");
578 : 5 : state->stash_tranche = LWLockNewTrancheId("pg_stash_advice_stash");
579 : 5 : state->entry_tranche = LWLockNewTrancheId("pg_stash_advice_entry");
580 : 5 : state->next_stash_id = UINT64CONST(1);
581 : 5 : state->area = DSA_HANDLE_INVALID;
582 : 5 : state->stash_hash = DSHASH_HANDLE_INVALID;
583 : 5 : state->entry_hash = DSHASH_HANDLE_INVALID;
584 : 5 : state->bgworker_pid = InvalidPid;
585 : 5 : pg_atomic_init_flag(&state->stashes_ready);
586 : 5 : pg_atomic_init_u64(&state->change_count, 0);
587 : :
588 : : /*
589 : : * If this module was loaded via shared_preload_libraries, then
590 : : * pg_stash_advice_persist is a GUC variable. If it's true, that means
591 : : * that we should lock out manual stash modifications until the dump file
592 : : * has been successfully loaded. If it's false, there's nothing to load,
593 : : * so we set stashes_ready immediately.
594 : : *
595 : : * If this module was not loaded via shared_preload_libraries, then
596 : : * pg_stash_advice_persist is not a GUC variable, but it will be false,
597 : : * which leads to the correct behavior.
598 : : */
599 [ + + ]: 5 : if (!pg_stash_advice_persist)
600 : 1 : pg_atomic_test_set_flag(&state->stashes_ready);
601 : 5 : }
602 : :
603 : : /*
604 : : * Check whether a string looks like a valid identifier. It must contain only
605 : : * ASCII identifier characters, and must not begin with a digit.
606 : : */
607 : : static bool
608 : 38 : pgsa_is_identifier(char *str)
609 : : {
610 [ + + + + ]: 38 : if (*str >= '0' && *str <= '9')
611 : 1 : return false;
612 : :
613 [ + + ]: 397 : while (*str != '\0')
614 : : {
615 : 361 : char c = *str++;
616 : :
617 [ + + + - : 361 : if ((c < '0' || c > '9') && (c < 'a' || c > 'z') &&
+ + - + +
+ ]
618 [ + - + + ]: 39 : (c < 'A' || c > 'Z') && c != '_')
619 : 1 : return false;
620 : : }
621 : :
622 : 36 : return true;
623 : : }
624 : :
625 : : /*
626 : : * Look up the integer ID that corresponds to the given stash name.
627 : : *
628 : : * Returns 0 if no such stash exists.
629 : : */
630 : : uint64
631 : 50 : pgsa_lookup_stash_id(char *stash_name)
632 : : {
633 : : pgsa_stash *stash;
634 : : uint64 stash_id;
635 : :
636 : : /* Search the shared hash table. */
637 : 50 : stash = dshash_find(pgsa_stash_dshash, stash_name, false);
638 [ + + ]: 50 : if (stash == NULL)
639 : 4 : return 0;
640 : 46 : stash_id = stash->pgsa_stash_id;
641 : 46 : dshash_release_lock(pgsa_stash_dshash, stash);
642 : :
643 : 46 : return stash_id;
644 : : }
645 : :
646 : : /*
647 : : * Store a new or updated advice string for the given advice stash and query ID.
648 : : */
649 : : void
650 : 14 : pgsa_set_advice_string(char *stash_name, int64 queryId, char *advice_string)
651 : : {
652 : : pgsa_entry *entry;
653 : : bool found;
654 : : pgsa_entry_key key;
655 : : uint64 stash_id;
656 : : dsa_pointer new_dp;
657 : : dsa_pointer old_dp;
658 : :
659 : : /*
660 : : * The caller must hold our lock, at least in shared mode. This is
661 : : * important for two reasons.
662 : : *
663 : : * First, it holds off interrupts, so that we can't bail out of this code
664 : : * after allocating DSA memory for the advice string and before storing
665 : : * the resulting pointer somewhere that others can find it.
666 : : *
667 : : * Second, we need to avoid a race against pgsa_drop_stash(). That
668 : : * function removes a stash_name->stash_id mapping and all the entries for
669 : : * that stash_id. Without the lock, there's a race condition no matter
670 : : * which of those things it does first, because as soon as we've looked up
671 : : * the stash ID, that whole function can execute before we do the rest of
672 : : * our work, which would result in us adding an entry for a stash that no
673 : : * longer exists.
674 : : */
675 : : Assert(LWLockHeldByMe(&pgsa_state->lock));
676 : :
677 : : /* Look up the stash ID. */
678 [ + + ]: 14 : if ((stash_id = pgsa_lookup_stash_id(stash_name)) == 0)
679 [ + - ]: 1 : ereport(ERROR,
680 : : errcode(ERRCODE_INVALID_PARAMETER_VALUE),
681 : : errmsg("advice stash \"%s\" does not exist", stash_name));
682 : :
683 : : /* Allocate space for the advice string. */
684 : 13 : new_dp = dsa_allocate(pgsa_dsa_area, strlen(advice_string) + 1);
685 : 13 : strcpy(dsa_get_address(pgsa_dsa_area, new_dp), advice_string);
686 : :
687 : : /* Attempt to insert an entry into the hash table. */
688 : 13 : memset(&key, 0, sizeof(pgsa_entry_key));
689 : 13 : key.pgsa_stash_id = stash_id;
690 : 13 : key.queryId = queryId;
691 : 13 : entry = dshash_find_or_insert_extended(pgsa_entry_dshash, &key, &found,
692 : : DSHASH_INSERT_NO_OOM);
693 : :
694 : : /*
695 : : * If it didn't work, bail out, being careful to free the shared memory
696 : : * we've already allocated before, since error cleanup will not do so.
697 : : */
698 [ - + ]: 13 : if (entry == NULL)
699 : : {
700 : 0 : dsa_free(pgsa_dsa_area, new_dp);
701 [ # # ]: 0 : ereport(ERROR,
702 : : errcode(ERRCODE_OUT_OF_MEMORY),
703 : : errmsg("out of memory"),
704 : : errdetail("Could not insert advice string into shared hash table."));
705 : : }
706 : :
707 : : /* Update the entry and release the lock. */
708 [ + + ]: 13 : old_dp = found ? entry->advice_string : InvalidDsaPointer;
709 : 13 : entry->advice_string = new_dp;
710 : 13 : dshash_release_lock(pgsa_entry_dshash, entry);
711 : :
712 : : /*
713 : : * We're not safe from leaks yet!
714 : : *
715 : : * There's now a pointer to new_dp in the entry that we just updated, but
716 : : * that means that there's no longer anything pointing to old_dp.
717 : : */
718 [ + + ]: 13 : if (DsaPointerIsValid(old_dp))
719 : 2 : dsa_free(pgsa_dsa_area, old_dp);
720 : :
721 : : /* Bump change count. */
722 : 13 : pg_atomic_add_fetch_u64(&pgsa_state->change_count, 1);
723 : 13 : }
724 : :
725 : : /*
726 : : * Start our worker process.
727 : : */
728 : : void
729 : 4 : pgsa_start_worker(void)
730 : : {
731 : 4 : BackgroundWorker worker = {0};
732 : : BackgroundWorkerHandle *handle;
733 : : BgwHandleStatus status;
734 : : pid_t pid;
735 : :
736 : 4 : worker.bgw_flags = BGWORKER_SHMEM_ACCESS;
737 : 4 : worker.bgw_start_time = BgWorkerStart_ConsistentState;
738 : 4 : worker.bgw_restart_time = BGW_DEFAULT_RESTART_INTERVAL;
739 : 4 : strcpy(worker.bgw_library_name, "pg_stash_advice");
740 : 4 : strcpy(worker.bgw_function_name, "pg_stash_advice_worker_main");
741 : 4 : strcpy(worker.bgw_name, "pg_stash_advice worker");
742 : 4 : strcpy(worker.bgw_type, "pg_stash_advice worker");
743 : :
744 : : /*
745 : : * If process_shared_preload_libraries_in_progress = true, we may be in
746 : : * the postmaster, in which case this will really register the worker, or
747 : : * we may be in a child process in an EXEC_BACKEND build, in which case it
748 : : * will silently do nothing (which is the correct behavior).
749 : : */
750 [ + - ]: 4 : if (process_shared_preload_libraries_in_progress)
751 : : {
752 : 4 : RegisterBackgroundWorker(&worker);
753 : 4 : return;
754 : : }
755 : :
756 : : /*
757 : : * If process_shared_preload_libraries_in_progress = false, we're being
758 : : * asked to start the worker after system startup time. In other words,
759 : : * unless this is single-user mode, we're not in the postmaster, so we
760 : : * should use RegisterDynamicBackgroundWorker and then wait for startup to
761 : : * complete. (If we do happen to be in single-user mode, this will error
762 : : * out, which is fine.)
763 : : */
764 : 0 : worker.bgw_notify_pid = MyProcPid;
765 [ # # ]: 0 : if (!RegisterDynamicBackgroundWorker(&worker, &handle))
766 [ # # ]: 0 : ereport(ERROR,
767 : : (errcode(ERRCODE_INSUFFICIENT_RESOURCES),
768 : : errmsg("could not register background process"),
769 : : errhint("You may need to increase \"max_worker_processes\".")));
770 : 0 : status = WaitForBackgroundWorkerStartup(handle, &pid);
771 [ # # ]: 0 : if (status != BGWH_STARTED)
772 [ # # ]: 0 : ereport(ERROR,
773 : : (errcode(ERRCODE_INSUFFICIENT_RESOURCES),
774 : : errmsg("could not start background process"),
775 : : errhint("More details may be available in the server log.")));
776 : : }
|
