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