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 : }
|
