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 "nodes/queryjumble.h"
17 : #include "pg_plan_advice.h"
18 : #include "pg_stash_advice.h"
19 : #include "storage/dsm_registry.h"
20 : #include "utils/guc.h"
21 : #include "utils/memutils.h"
22 :
23 2 : PG_MODULE_MAGIC;
24 :
25 : /* Shared memory hash table parameters */
26 : static dshash_parameters pgsa_stash_dshash_parameters = {
27 : NAMEDATALEN,
28 : sizeof(pgsa_stash),
29 : dshash_strcmp,
30 : dshash_strhash,
31 : dshash_strcpy,
32 : LWTRANCHE_INVALID /* gets set at runtime */
33 : };
34 :
35 : static dshash_parameters pgsa_entry_dshash_parameters = {
36 : sizeof(pgsa_entry_key),
37 : sizeof(pgsa_entry),
38 : dshash_memcmp,
39 : dshash_memhash,
40 : dshash_memcpy,
41 : LWTRANCHE_INVALID /* gets set at runtime */
42 : };
43 :
44 : /* GUC variable */
45 : static char *pg_stash_advice_stash_name = "";
46 :
47 : /* Shared memory pointers */
48 : pgsa_shared_state *pgsa_state;
49 : dsa_area *pgsa_dsa_area;
50 : dshash_table *pgsa_stash_dshash;
51 : dshash_table *pgsa_entry_dshash;
52 :
53 : /* Other global variables */
54 : static MemoryContext pg_stash_advice_mcxt;
55 :
56 : /* Function prototypes */
57 : static char *pgsa_advisor(PlannerGlobal *glob,
58 : Query *parse,
59 : const char *query_string,
60 : int cursorOptions,
61 : ExplainState *es);
62 : static bool pgsa_check_stash_name_guc(char **newval, void **extra,
63 : GucSource source);
64 : static void pgsa_init_shared_state(void *ptr, void *arg);
65 : static bool pgsa_is_identifier(char *str);
66 :
67 : /* Stash name -> stash ID hash table */
68 : #define SH_PREFIX pgsa_stash_name_table
69 : #define SH_ELEMENT_TYPE pgsa_stash_name
70 : #define SH_KEY_TYPE uint64
71 : #define SH_KEY pgsa_stash_id
72 : #define SH_HASH_KEY(tb, key) hash_bytes((const unsigned char *) &(key), sizeof(uint64))
73 : #define SH_EQUAL(tb, a, b) (a == b)
74 : #define SH_SCOPE extern
75 : #define SH_DEFINE
76 : #include "lib/simplehash.h"
77 :
78 : /*
79 : * Initialize this module.
80 : */
81 : void
82 2 : _PG_init(void)
83 : {
84 : void (*add_advisor_fn) (pg_plan_advice_advisor_hook hook);
85 :
86 : /* If compute_query_id = 'auto', we would like query IDs. */
87 2 : EnableQueryId();
88 :
89 : /* Define our GUCs. */
90 2 : DefineCustomStringVariable("pg_stash_advice.stash_name",
91 : "Name of the advice stash to be used in this session.",
92 : NULL,
93 : &pg_stash_advice_stash_name,
94 : "",
95 : PGC_USERSET,
96 : 0,
97 : pgsa_check_stash_name_guc,
98 : NULL,
99 : NULL);
100 :
101 2 : MarkGUCPrefixReserved("pg_stash_advice");
102 :
103 : /* Tell pg_plan_advice that we want to provide advice strings. */
104 2 : add_advisor_fn =
105 2 : load_external_function("pg_plan_advice", "pg_plan_advice_add_advisor",
106 : true, NULL);
107 2 : (*add_advisor_fn) (pgsa_advisor);
108 2 : }
109 :
110 : /*
111 : * Get the advice string that has been configured for this query, if any,
112 : * and return it. Otherwise, return NULL.
113 : */
114 : static char *
115 39 : pgsa_advisor(PlannerGlobal *glob, Query *parse,
116 : const char *query_string, int cursorOptions,
117 : ExplainState *es)
118 : {
119 : pgsa_entry_key key;
120 : pgsa_entry *entry;
121 : char *advice_string;
122 : uint64 stash_id;
123 :
124 : /*
125 : * Exit quickly if the stash name is empty or there's no query ID.
126 : */
127 39 : if (pg_stash_advice_stash_name[0] == '\0' || parse->queryId == 0)
128 9 : return NULL;
129 :
130 : /* Attach to dynamic shared memory if not already done. */
131 30 : if (unlikely(pgsa_entry_dshash == NULL))
132 0 : pgsa_attach();
133 :
134 : /*
135 : * Translate pg_stash_advice.stash_name to an integer ID.
136 : *
137 : * pgsa_check_stash_name_guc() has already validated the advice stash
138 : * name, so we don't need to call pgsa_check_stash_name() here.
139 : */
140 30 : stash_id = pgsa_lookup_stash_id(pg_stash_advice_stash_name);
141 30 : if (stash_id == 0)
142 1 : return NULL;
143 :
144 : /*
145 : * Look up the advice string for the given stash ID + query ID.
146 : *
147 : * If we find an advice string, we copy it into the current memory
148 : * context, presumably short-lived, so that we can release the lock on the
149 : * dshash entry. pg_plan_advice only needs the value to remain allocated
150 : * long enough for it to be parsed, so this should be good enough.
151 : */
152 29 : memset(&key, 0, sizeof(pgsa_entry_key));
153 29 : key.pgsa_stash_id = stash_id;
154 29 : key.queryId = parse->queryId;
155 29 : entry = dshash_find(pgsa_entry_dshash, &key, false);
156 29 : if (entry == NULL)
157 25 : return NULL;
158 4 : if (entry->advice_string == InvalidDsaPointer)
159 0 : advice_string = NULL;
160 : else
161 4 : advice_string = pstrdup(dsa_get_address(pgsa_dsa_area,
162 : entry->advice_string));
163 4 : dshash_release_lock(pgsa_entry_dshash, entry);
164 :
165 : /* If we found an advice string, emit a debug message. */
166 4 : if (advice_string != NULL)
167 4 : elog(DEBUG2, "supplying automatic advice for stash \"%s\", query ID %" PRId64 ": %s",
168 : pg_stash_advice_stash_name, key.queryId, advice_string);
169 :
170 4 : return advice_string;
171 : }
172 :
173 : /*
174 : * Attach to various structures in dynamic shared memory.
175 : *
176 : * This function is designed to be resilient against errors. That is, if it
177 : * fails partway through, it should be possible to call it again, repeat no
178 : * work already completed, and potentially succeed or at least get further if
179 : * whatever caused the previous failure has been corrected.
180 : */
181 : void
182 1 : pgsa_attach(void)
183 : {
184 : bool found;
185 : MemoryContext oldcontext;
186 :
187 : /*
188 : * Create a memory context to make sure that any control structures
189 : * allocated in local memory are sufficiently persistent.
190 : */
191 1 : if (pg_stash_advice_mcxt == NULL)
192 1 : pg_stash_advice_mcxt = AllocSetContextCreate(TopMemoryContext,
193 : "pg_stash_advice",
194 : ALLOCSET_DEFAULT_SIZES);
195 1 : oldcontext = MemoryContextSwitchTo(pg_stash_advice_mcxt);
196 :
197 : /* Attach to the fixed-size state object if not already done. */
198 1 : if (pgsa_state == NULL)
199 1 : pgsa_state = GetNamedDSMSegment("pg_stash_advice",
200 : sizeof(pgsa_shared_state),
201 : pgsa_init_shared_state,
202 : &found, NULL);
203 :
204 : /* Attach to the DSA area if not already done. */
205 1 : if (pgsa_dsa_area == NULL)
206 : {
207 : dsa_handle area_handle;
208 :
209 1 : LWLockAcquire(&pgsa_state->lock, LW_EXCLUSIVE);
210 1 : area_handle = pgsa_state->area;
211 1 : if (area_handle == DSA_HANDLE_INVALID)
212 : {
213 1 : pgsa_dsa_area = dsa_create(pgsa_state->dsa_tranche);
214 1 : dsa_pin(pgsa_dsa_area);
215 1 : pgsa_state->area = dsa_get_handle(pgsa_dsa_area);
216 1 : LWLockRelease(&pgsa_state->lock);
217 : }
218 : else
219 : {
220 0 : LWLockRelease(&pgsa_state->lock);
221 0 : pgsa_dsa_area = dsa_attach(area_handle);
222 : }
223 1 : dsa_pin_mapping(pgsa_dsa_area);
224 : }
225 :
226 : /* Attach to the stash_name->stash_id hash table if not already done. */
227 1 : if (pgsa_stash_dshash == NULL)
228 : {
229 : dshash_table_handle stash_handle;
230 :
231 1 : LWLockAcquire(&pgsa_state->lock, LW_EXCLUSIVE);
232 1 : pgsa_stash_dshash_parameters.tranche_id = pgsa_state->stash_tranche;
233 1 : stash_handle = pgsa_state->stash_hash;
234 1 : if (stash_handle == DSHASH_HANDLE_INVALID)
235 : {
236 1 : pgsa_stash_dshash = dshash_create(pgsa_dsa_area,
237 : &pgsa_stash_dshash_parameters,
238 : NULL);
239 2 : pgsa_state->stash_hash =
240 1 : dshash_get_hash_table_handle(pgsa_stash_dshash);
241 1 : LWLockRelease(&pgsa_state->lock);
242 : }
243 : else
244 : {
245 0 : LWLockRelease(&pgsa_state->lock);
246 0 : pgsa_stash_dshash = dshash_attach(pgsa_dsa_area,
247 : &pgsa_stash_dshash_parameters,
248 : stash_handle, NULL);
249 : }
250 : }
251 :
252 : /* Attach to the entry hash table if not already done. */
253 1 : if (pgsa_entry_dshash == NULL)
254 : {
255 : dshash_table_handle entry_handle;
256 :
257 1 : LWLockAcquire(&pgsa_state->lock, LW_EXCLUSIVE);
258 1 : pgsa_entry_dshash_parameters.tranche_id = pgsa_state->entry_tranche;
259 1 : entry_handle = pgsa_state->entry_hash;
260 1 : if (entry_handle == DSHASH_HANDLE_INVALID)
261 : {
262 1 : pgsa_entry_dshash = dshash_create(pgsa_dsa_area,
263 : &pgsa_entry_dshash_parameters,
264 : NULL);
265 2 : pgsa_state->entry_hash =
266 1 : dshash_get_hash_table_handle(pgsa_entry_dshash);
267 1 : LWLockRelease(&pgsa_state->lock);
268 : }
269 : else
270 : {
271 0 : LWLockRelease(&pgsa_state->lock);
272 0 : pgsa_entry_dshash = dshash_attach(pgsa_dsa_area,
273 : &pgsa_entry_dshash_parameters,
274 : entry_handle, NULL);
275 : }
276 : }
277 :
278 : /* Restore previous memory context. */
279 1 : MemoryContextSwitchTo(oldcontext);
280 1 : }
281 :
282 : /*
283 : * Check whether an advice stash name is legal, and signal an error if not.
284 : *
285 : * Keep this in sync with pgsa_check_stash_name_guc, below.
286 : */
287 : void
288 22 : pgsa_check_stash_name(char *stash_name)
289 : {
290 : /* Reject empty advice stash name. */
291 22 : if (stash_name[0] == '\0')
292 1 : ereport(ERROR,
293 : errcode(ERRCODE_INVALID_PARAMETER_VALUE),
294 : errmsg("advice stash name may not be zero length"));
295 :
296 : /* Reject overlong advice stash names. */
297 21 : if (strlen(stash_name) + 1 > NAMEDATALEN)
298 1 : ereport(ERROR,
299 : errcode(ERRCODE_INVALID_PARAMETER_VALUE),
300 : errmsg("advice stash names may not be longer than %d bytes",
301 : NAMEDATALEN - 1));
302 :
303 : /*
304 : * Reject non-ASCII advice stash names, since advice stashes are visible
305 : * across all databases and the encodings of those databases might differ.
306 : */
307 20 : if (!pg_is_ascii(stash_name))
308 1 : ereport(ERROR,
309 : errcode(ERRCODE_INVALID_PARAMETER_VALUE),
310 : errmsg("advice stash name must not contain non-ASCII characters"));
311 :
312 : /*
313 : * Reject things that do not look like identifiers, since the ability to
314 : * create an advice stash with non-printable characters or weird symbols
315 : * in the name is not likely to be useful to anyone.
316 : */
317 19 : if (!pgsa_is_identifier(stash_name))
318 1 : ereport(ERROR,
319 : errcode(ERRCODE_INVALID_PARAMETER_VALUE),
320 : errmsg("advice stash name must begin with a letter or underscore and contain only letters, digits, and underscores"));
321 18 : }
322 :
323 : /*
324 : * As above, but for the GUC check_hook. We allow the empty string here,
325 : * though, as equivalent to disabling the feature.
326 : */
327 : static bool
328 7 : pgsa_check_stash_name_guc(char **newval, void **extra, GucSource source)
329 : {
330 7 : char *stash_name = *newval;
331 :
332 : /* Reject overlong advice stash names. */
333 7 : if (strlen(stash_name) + 1 > NAMEDATALEN)
334 : {
335 0 : GUC_check_errcode(ERRCODE_INVALID_PARAMETER_VALUE);
336 0 : GUC_check_errdetail("advice stash names may not be longer than %d bytes",
337 : NAMEDATALEN - 1);
338 0 : return false;
339 : }
340 :
341 : /*
342 : * Reject non-ASCII advice stash names, since advice stashes are visible
343 : * across all databases and the encodings of those databases might differ.
344 : */
345 7 : if (!pg_is_ascii(stash_name))
346 : {
347 1 : GUC_check_errcode(ERRCODE_INVALID_PARAMETER_VALUE);
348 1 : GUC_check_errdetail("advice stash name must not contain non-ASCII characters");
349 1 : return false;
350 : }
351 :
352 : /*
353 : * Reject things that do not look like identifiers, since the ability to
354 : * create an advice stash with non-printable characters or weird symbols
355 : * in the name is not likely to be useful to anyone.
356 : */
357 6 : if (!pgsa_is_identifier(stash_name))
358 : {
359 1 : GUC_check_errcode(ERRCODE_INVALID_PARAMETER_VALUE);
360 1 : GUC_check_errdetail("advice stash name must begin with a letter or underscore and contain only letters, digits, and underscores");
361 1 : return false;
362 : }
363 :
364 5 : return true;
365 : }
366 :
367 : /*
368 : * Create an advice stash.
369 : */
370 : void
371 3 : pgsa_create_stash(char *stash_name)
372 : {
373 : pgsa_stash *stash;
374 : bool found;
375 :
376 : Assert(LWLockHeldByMeInMode(&pgsa_state->lock, LW_EXCLUSIVE));
377 :
378 : /* Create a stash with this name, unless one already exists. */
379 3 : stash = dshash_find_or_insert(pgsa_stash_dshash, stash_name, &found);
380 3 : if (found)
381 1 : ereport(ERROR,
382 : errcode(ERRCODE_INVALID_PARAMETER_VALUE),
383 : errmsg("advice stash \"%s\" already exists", stash_name));
384 2 : stash->pgsa_stash_id = pgsa_state->next_stash_id++;
385 2 : dshash_release_lock(pgsa_stash_dshash, stash);
386 2 : }
387 :
388 : /*
389 : * Remove any stored advice string for the given advice stash and query ID.
390 : */
391 : void
392 2 : pgsa_clear_advice_string(char *stash_name, int64 queryId)
393 : {
394 : pgsa_entry *entry;
395 : pgsa_entry_key key;
396 : uint64 stash_id;
397 : dsa_pointer old_dp;
398 :
399 : Assert(LWLockHeldByMe(&pgsa_state->lock));
400 :
401 : /* Translate the stash name to an integer ID. */
402 2 : if ((stash_id = pgsa_lookup_stash_id(stash_name)) == 0)
403 1 : ereport(ERROR,
404 : errcode(ERRCODE_INVALID_PARAMETER_VALUE),
405 : errmsg("advice stash \"%s\" does not exist", stash_name));
406 :
407 : /*
408 : * Look for an existing entry, and free it. But, be sure to save the
409 : * pointer to the associated advice string, if any.
410 : */
411 1 : memset(&key, 0, sizeof(pgsa_entry_key));
412 1 : key.pgsa_stash_id = stash_id;
413 1 : key.queryId = queryId;
414 1 : entry = dshash_find(pgsa_entry_dshash, &key, true);
415 1 : if (entry == NULL)
416 0 : old_dp = InvalidDsaPointer;
417 : else
418 : {
419 1 : old_dp = entry->advice_string;
420 1 : dshash_delete_entry(pgsa_entry_dshash, entry);
421 : }
422 :
423 : /* Now we free the advice string as well, if there was one. */
424 1 : if (old_dp != InvalidDsaPointer)
425 1 : dsa_free(pgsa_dsa_area, old_dp);
426 1 : }
427 :
428 : /*
429 : * Drop an advice stash.
430 : */
431 : void
432 3 : pgsa_drop_stash(char *stash_name)
433 : {
434 : pgsa_entry *entry;
435 : pgsa_stash *stash;
436 : dshash_seq_status iterator;
437 : uint64 stash_id;
438 :
439 : Assert(LWLockHeldByMeInMode(&pgsa_state->lock, LW_EXCLUSIVE));
440 :
441 : /* Remove the entry for this advice stash. */
442 3 : stash = dshash_find(pgsa_stash_dshash, stash_name, true);
443 3 : if (stash == NULL)
444 1 : ereport(ERROR,
445 : errcode(ERRCODE_INVALID_PARAMETER_VALUE),
446 : errmsg("advice stash \"%s\" does not exist", stash_name));
447 2 : stash_id = stash->pgsa_stash_id;
448 2 : dshash_delete_entry(pgsa_stash_dshash, stash);
449 :
450 : /*
451 : * Now remove all the entries. Since pgsa_state->lock must be held at
452 : * least in shared mode to insert entries into pgsa_entry_dshash, it
453 : * doesn't matter whether we do this before or after deleting the entry
454 : * from pgsa_stash_dshash.
455 : */
456 2 : dshash_seq_init(&iterator, pgsa_entry_dshash, true);
457 5 : while ((entry = dshash_seq_next(&iterator)) != NULL)
458 : {
459 1 : if (stash_id == entry->key.pgsa_stash_id)
460 : {
461 1 : if (entry->advice_string != InvalidDsaPointer)
462 1 : dsa_free(pgsa_dsa_area, entry->advice_string);
463 1 : dshash_delete_current(&iterator);
464 : }
465 : }
466 2 : dshash_seq_term(&iterator);
467 2 : }
468 :
469 : /*
470 : * Initialize shared state when first created.
471 : */
472 : static void
473 1 : pgsa_init_shared_state(void *ptr, void *arg)
474 : {
475 1 : pgsa_shared_state *state = (pgsa_shared_state *) ptr;
476 :
477 1 : LWLockInitialize(&state->lock,
478 : LWLockNewTrancheId("pg_stash_advice_lock"));
479 1 : state->dsa_tranche = LWLockNewTrancheId("pg_stash_advice_dsa");
480 1 : state->stash_tranche = LWLockNewTrancheId("pg_stash_advice_stash");
481 1 : state->entry_tranche = LWLockNewTrancheId("pg_stash_advice_entry");
482 1 : state->next_stash_id = UINT64CONST(1);
483 1 : state->area = DSA_HANDLE_INVALID;
484 1 : state->stash_hash = DSHASH_HANDLE_INVALID;
485 1 : state->entry_hash = DSHASH_HANDLE_INVALID;
486 1 : }
487 :
488 : /*
489 : * Check whether a string looks like a valid identifier. It must contain only
490 : * ASCII identifier characters, and must not begin with a digit.
491 : */
492 : static bool
493 25 : pgsa_is_identifier(char *str)
494 : {
495 25 : if (*str >= '0' && *str <= '9')
496 1 : return false;
497 :
498 321 : while (*str != '\0')
499 : {
500 298 : char c = *str++;
501 :
502 298 : if ((c < '0' || c > '9') && (c < 'a' || c > 'z') &&
503 30 : (c < 'A' || c > 'Z') && c != '_')
504 1 : return false;
505 : }
506 :
507 23 : return true;
508 : }
509 :
510 : /*
511 : * Look up the integer ID that corresponds to the given stash name.
512 : *
513 : * Returns 0 if no such stash exists.
514 : */
515 : uint64
516 41 : pgsa_lookup_stash_id(char *stash_name)
517 : {
518 : pgsa_stash *stash;
519 : uint64 stash_id;
520 :
521 : /* Search the shared hash table. */
522 41 : stash = dshash_find(pgsa_stash_dshash, stash_name, false);
523 41 : if (stash == NULL)
524 4 : return 0;
525 37 : stash_id = stash->pgsa_stash_id;
526 37 : dshash_release_lock(pgsa_stash_dshash, stash);
527 :
528 37 : return stash_id;
529 : }
530 :
531 : /*
532 : * Store a new or updated advice string for the given advice stash and query ID.
533 : */
534 : void
535 5 : pgsa_set_advice_string(char *stash_name, int64 queryId, char *advice_string)
536 : {
537 : pgsa_entry *entry;
538 : bool found;
539 : pgsa_entry_key key;
540 : uint64 stash_id;
541 : dsa_pointer new_dp;
542 : dsa_pointer old_dp;
543 :
544 : /*
545 : * The caller must hold our lock, at least in shared mode. This is
546 : * important for two reasons.
547 : *
548 : * First, it holds off interrupts, so that we can't bail out of this code
549 : * after allocating DSA memory for the advice string and before storing
550 : * the resulting pointer somewhere that others can find it.
551 : *
552 : * Second, we need to avoid a race against pgsa_drop_stash(). That
553 : * function removes a stash_name->stash_id mapping and all the entries for
554 : * that stash_id. Without the lock, there's a race condition no matter
555 : * which of those things it does first, because as soon as we've looked up
556 : * the stash ID, that whole function can execute before we do the rest of
557 : * our work, which would result in us adding an entry for a stash that no
558 : * longer exists.
559 : */
560 : Assert(LWLockHeldByMe(&pgsa_state->lock));
561 :
562 : /* Look up the stash ID. */
563 5 : if ((stash_id = pgsa_lookup_stash_id(stash_name)) == 0)
564 1 : ereport(ERROR,
565 : errcode(ERRCODE_INVALID_PARAMETER_VALUE),
566 : errmsg("advice stash \"%s\" does not exist", stash_name));
567 :
568 : /* Allocate space for the advice string. */
569 4 : new_dp = dsa_allocate(pgsa_dsa_area, strlen(advice_string) + 1);
570 4 : strcpy(dsa_get_address(pgsa_dsa_area, new_dp), advice_string);
571 :
572 : /* Attempt to insert an entry into the hash table. */
573 4 : memset(&key, 0, sizeof(pgsa_entry_key));
574 4 : key.pgsa_stash_id = stash_id;
575 4 : key.queryId = queryId;
576 4 : entry = dshash_find_or_insert_extended(pgsa_entry_dshash, &key, &found,
577 : DSHASH_INSERT_NO_OOM);
578 :
579 : /*
580 : * If it didn't work, bail out, being careful to free the shared memory
581 : * we've already allocated before, since error cleanup will not do so.
582 : */
583 4 : if (entry == NULL)
584 : {
585 0 : dsa_free(pgsa_dsa_area, new_dp);
586 0 : ereport(ERROR,
587 : errcode(ERRCODE_OUT_OF_MEMORY),
588 : errmsg("out of memory"),
589 : errdetail("could not insert advice string into shared hash table"));
590 : }
591 :
592 : /* Update the entry and release the lock. */
593 4 : old_dp = found ? entry->advice_string : InvalidDsaPointer;
594 4 : entry->advice_string = new_dp;
595 4 : dshash_release_lock(pgsa_entry_dshash, entry);
596 :
597 : /*
598 : * We're not safe from leaks yet!
599 : *
600 : * There's now a pointer to new_dp in the entry that we just updated, but
601 : * that means that there's no longer anything pointing to old_dp.
602 : */
603 4 : if (DsaPointerIsValid(old_dp))
604 2 : dsa_free(pgsa_dsa_area, old_dp);
605 4 : }
|
