postgresql/contrib/pg_stash_advice/stashfuncs.c

/*-------------------------------------------------------------------------
 *
 * stashfuncs.c
 *	  SQL interface to pg_stash_advice
 *
 * Copyright (c) 2016-2026, PostgreSQL Global Development Group
 *
 *	  contrib/pg_stash_advice/stashfuncs.c
 *
 *-------------------------------------------------------------------------
 */
#include "postgres.h"

#include "common/hashfn.h"
#include "fmgr.h"
#include "funcapi.h"
#include "miscadmin.h"
#include "pg_stash_advice.h"
#include "utils/builtins.h"
#include "utils/tuplestore.h"

PG_FUNCTION_INFO_V1(pg_create_advice_stash);
PG_FUNCTION_INFO_V1(pg_drop_advice_stash);
PG_FUNCTION_INFO_V1(pg_get_advice_stash_contents);
PG_FUNCTION_INFO_V1(pg_get_advice_stashes);
PG_FUNCTION_INFO_V1(pg_set_stashed_advice);
PG_FUNCTION_INFO_V1(pg_start_stash_advice_worker);

typedef struct pgsa_stash_count
{
	uint32		status;
	uint64		pgsa_stash_id;
	int64		num_entries;
} pgsa_stash_count;

#define SH_PREFIX pgsa_stash_count_table
#define SH_ELEMENT_TYPE pgsa_stash_count
#define SH_KEY_TYPE uint64
#define SH_KEY pgsa_stash_id
#define SH_HASH_KEY(tb, key) hash_bytes((const unsigned char *) &(key), sizeof(uint64))
#define SH_EQUAL(tb, a, b) (a == b)
#define SH_SCOPE static inline
#define SH_DEFINE
#define SH_DECLARE
#include "lib/simplehash.h"

/*
 * SQL-callable function to create an advice stash
 */
Datum
pg_create_advice_stash(PG_FUNCTION_ARGS)
{
	char	   *stash_name = text_to_cstring(PG_GETARG_TEXT_PP(0));

	pgsa_check_stash_name(stash_name);
	if (unlikely(pgsa_entry_dshash == NULL))
		pgsa_attach();
	pgsa_check_lockout();
	LWLockAcquire(&pgsa_state->lock, LW_EXCLUSIVE);
	pgsa_create_stash(stash_name);
	LWLockRelease(&pgsa_state->lock);
	PG_RETURN_VOID();
}

/*
 * SQL-callable function to drop an advice stash
 */
Datum
pg_drop_advice_stash(PG_FUNCTION_ARGS)
{
	char	   *stash_name = text_to_cstring(PG_GETARG_TEXT_PP(0));

	pgsa_check_stash_name(stash_name);
	if (unlikely(pgsa_entry_dshash == NULL))
		pgsa_attach();
	pgsa_check_lockout();
	LWLockAcquire(&pgsa_state->lock, LW_EXCLUSIVE);
	pgsa_drop_stash(stash_name);
	LWLockRelease(&pgsa_state->lock);
	PG_RETURN_VOID();
}

/*
 * SQL-callable function to provide a list of advice stashes
 */
Datum
pg_get_advice_stashes(PG_FUNCTION_ARGS)
{
	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
	dshash_seq_status iterator;
	pgsa_entry *entry;
	pgsa_stash *stash;
	pgsa_stash_count_table_hash *chash;

	InitMaterializedSRF(fcinfo, 0);

	/* Attach to dynamic shared memory if not already done. */
	if (unlikely(pgsa_entry_dshash == NULL))
		pgsa_attach();

	/* If stash data is still being restored from disk, ignore. */
	if (pg_atomic_unlocked_test_flag(&pgsa_state->stashes_ready))
		return (Datum) 0;

	/* Tally up the number of entries per stash. */
	chash = pgsa_stash_count_table_create(CurrentMemoryContext, 64, NULL);
	dshash_seq_init(&iterator, pgsa_entry_dshash, true);
	while ((entry = dshash_seq_next(&iterator)) != NULL)
	{
		pgsa_stash_count *c;
		bool		found;

		c = pgsa_stash_count_table_insert(chash,
										  entry->key.pgsa_stash_id,
										  &found);
		if (!found)
			c->num_entries = 1;
		else
			c->num_entries++;
	}
	dshash_seq_term(&iterator);

	/* Emit results. */
	dshash_seq_init(&iterator, pgsa_stash_dshash, true);
	while ((stash = dshash_seq_next(&iterator)) != NULL)
	{
		Datum		values[2];
		bool		nulls[2];
		pgsa_stash_count *c;

		values[0] = CStringGetTextDatum(stash->name);
		nulls[0] = false;

		c = pgsa_stash_count_table_lookup(chash, stash->pgsa_stash_id);
		values[1] = Int64GetDatum(c == NULL ? 0 : c->num_entries);
		nulls[1] = false;

		tuplestore_putvalues(rsinfo->setResult, rsinfo->setDesc, values,
							 nulls);
	}
	dshash_seq_term(&iterator);

	return (Datum) 0;
}

/*
 * SQL-callable function to provide advice stash contents
 */
Datum
pg_get_advice_stash_contents(PG_FUNCTION_ARGS)
{
	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
	dshash_seq_status iterator;
	char	   *stash_name = NULL;
	pgsa_stash_name_table_hash *nhash = NULL;
	uint64		stash_id = 0;
	pgsa_entry *entry;

	InitMaterializedSRF(fcinfo, 0);

	/* Attach to dynamic shared memory if not already done. */
	if (unlikely(pgsa_entry_dshash == NULL))
		pgsa_attach();

	/* If stash data is still being restored from disk, ignore. */
	if (pg_atomic_unlocked_test_flag(&pgsa_state->stashes_ready))
		return (Datum) 0;

	/* User can pass NULL for all stashes, or the name of a specific stash. */
	if (!PG_ARGISNULL(0))
	{
		stash_name = text_to_cstring(PG_GETARG_TEXT_PP(0));
		pgsa_check_stash_name(stash_name);
		stash_id = pgsa_lookup_stash_id(stash_name);

		/* If the user specified a stash name, it should exist. */
		if (stash_id == 0)
			ereport(ERROR,
					errcode(ERRCODE_INVALID_PARAMETER_VALUE),
					errmsg("advice stash \"%s\" does not exist", stash_name));
	}
	else
	{
		pgsa_stash *stash;

		/*
		 * If we're dumping data about all stashes, we need an ID->name lookup
		 * table.
		 */
		nhash = pgsa_stash_name_table_create(CurrentMemoryContext, 64, NULL);
		dshash_seq_init(&iterator, pgsa_stash_dshash, true);
		while ((stash = dshash_seq_next(&iterator)) != NULL)
		{
			pgsa_stash_name *n;
			bool		found;

			n = pgsa_stash_name_table_insert(nhash,
											 stash->pgsa_stash_id,
											 &found);
			Assert(!found);
			n->name = pstrdup(stash->name);
		}
		dshash_seq_term(&iterator);
	}

	/* Now iterate over all the entries. */
	dshash_seq_init(&iterator, pgsa_entry_dshash, false);
	while ((entry = dshash_seq_next(&iterator)) != NULL)
	{
		Datum		values[3];
		bool		nulls[3];
		char	   *this_stash_name;
		char	   *advice_string;

		/* Skip incomplete entries where the advice string was never set. */
		if (entry->advice_string == InvalidDsaPointer)
			continue;

		if (stash_id != 0)
		{
			/*
			 * We're only dumping data for one particular stash, so skip
			 * entries for any other stash and use the stash name specified by
			 * the user.
			 */
			if (stash_id != entry->key.pgsa_stash_id)
				continue;
			this_stash_name = stash_name;
		}
		else
		{
			pgsa_stash_name *n;

			/*
			 * We're dumping data for all stashes, so look up the correct name
			 * to use in the hash table. If nothing is found, which is
			 * possible due to race conditions, make up a string to use.
			 */
			n = pgsa_stash_name_table_lookup(nhash, entry->key.pgsa_stash_id);
			if (n != NULL)
				this_stash_name = n->name;
			else
				this_stash_name = psprintf("<stash %" PRIu64 ">",
										   entry->key.pgsa_stash_id);
		}

		/* Work out tuple values. */
		values[0] = CStringGetTextDatum(this_stash_name);
		nulls[0] = false;
		values[1] = Int64GetDatum(entry->key.queryId);
		nulls[1] = false;
		advice_string = dsa_get_address(pgsa_dsa_area, entry->advice_string);
		values[2] = CStringGetTextDatum(advice_string);
		nulls[2] = false;

		/* Emit the tuple. */
		tuplestore_putvalues(rsinfo->setResult, rsinfo->setDesc, values,
							 nulls);
	}
	dshash_seq_term(&iterator);

	return (Datum) 0;
}

/*
 * SQL-callable function to update an advice stash entry for a particular
 * query ID
 *
 * If the second argument is NULL, we delete any existing advice stash
 * entry; otherwise, we either create an entry or update it with the new
 * advice string.
 */
Datum
pg_set_stashed_advice(PG_FUNCTION_ARGS)
{
	char	   *stash_name;
	int64		queryId;

	if (PG_ARGISNULL(0) || PG_ARGISNULL(1))
		PG_RETURN_NULL();

	/* Get and check advice stash name. */
	stash_name = text_to_cstring(PG_GETARG_TEXT_PP(0));
	pgsa_check_stash_name(stash_name);

	/*
	 * Get and check query ID.
	 *
	 * Query ID 0 means no query ID was computed, so reject that.
	 */
	queryId = PG_GETARG_INT64(1);
	if (queryId == 0)
		ereport(ERROR,
				errcode(ERRCODE_INVALID_PARAMETER_VALUE),
				errmsg("cannot set advice string for query ID 0"));

	/* Attach to dynamic shared memory if not already done. */
	if (unlikely(pgsa_entry_dshash == NULL))
		pgsa_attach();

	/* Don't allow writes if stash data is still being restored from disk. */
	pgsa_check_lockout();

	/* Now call the appropriate function to do the real work. */
	if (PG_ARGISNULL(2))
	{
		LWLockAcquire(&pgsa_state->lock, LW_SHARED);
		pgsa_clear_advice_string(stash_name, queryId);
		LWLockRelease(&pgsa_state->lock);
	}
	else
	{
		char	   *advice_string = text_to_cstring(PG_GETARG_TEXT_PP(2));

		LWLockAcquire(&pgsa_state->lock, LW_SHARED);
		pgsa_set_advice_string(stash_name, queryId, advice_string);
		LWLockRelease(&pgsa_state->lock);
	}

	PG_RETURN_VOID();
}

/*
 * SQL-callable function to start the persistence background worker.
 */
Datum
pg_start_stash_advice_worker(PG_FUNCTION_ARGS)
{
	pid_t		pid;

	if (unlikely(pgsa_entry_dshash == NULL))
		pgsa_attach();

	LWLockAcquire(&pgsa_state->lock, LW_SHARED);
	pid = pgsa_state->bgworker_pid;
	LWLockRelease(&pgsa_state->lock);

	if (pid != InvalidPid)
		ereport(ERROR,
				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
				 errmsg("pg_stash_advice worker is already running under PID %d",
						(int) pid)));

	pgsa_start_worker();

	PG_RETURN_VOID();
}