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