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