/*--------------------------------------------------------------------------
 *
 * pg_stat_plans.c
 *		Track per-plan call counts, execution times and EXPLAIN texts
 *		across a whole database cluster.
 *
 * Portions Copyright (c) 1996-2024, PostgreSQL Global Development Group
 * Portions Copyright (c) 1994, Regents of the University of California
 *
 * IDENTIFICATION
 *		contrib/pg_stat_plans/pg_stat_plans.c
 *
 * -------------------------------------------------------------------------
 */

#include "postgres.h"

#include "access/htup_details.h"
#include "access/parallel.h"
#include "catalog/pg_authid.h"
#include "commands/explain.h"
#if PG_VERSION_NUM >= 180000
#include "commands/explain_format.h"
#include "commands/explain_state.h"
#endif
#include "common/hashfn.h"
#include "funcapi.h"
#include "lib/dshash.h"
#include "libpq/auth.h"
#include "mb/pg_wchar.h"
#include "nodes/parsenodes.h"
#include "nodes/queryjumble.h"
#include "optimizer/planner.h"
#include "parser/analyze.h"
#include "pgstat.h"
#if PG_VERSION_NUM >= 190000
#include "storage/dsm_registry.h"
#include "utils/pgstat_internal.h"
#endif
#include "storage/ipc.h"
#include "storage/lwlock.h"
#include "storage/shmem.h"
#if PG_VERSION_NUM >= 180000
#include "storage/dsm_registry.h"
#endif
#include "tcop/utility.h"
#include "utils/dsa.h"
#include "utils/acl.h"
#include "utils/builtins.h"
#include "utils/guc.h"
#include "utils/memutils.h"
#include "utils/snapmgr.h"
#include "utils/tuplestore.h"

#include "jumblefuncs.h"

#include "compat_16_17/pgstat_custom.h"

#ifdef USE_ZSTD
#include <zstd.h>
#endif

PG_MODULE_MAGIC;

/* Current nesting depth of planner/ExecutorRun/ProcessUtility calls */
static int	nesting_level = 0;

/* Saved hook values */
static post_parse_analyze_hook_type prev_post_parse_analyze_hook = NULL;
static planner_hook_type prev_planner_hook = NULL;
static ExecutorStart_hook_type prev_ExecutorStart = NULL;
static ExecutorRun_hook_type prev_ExecutorRun = NULL;
static ExecutorFinish_hook_type prev_ExecutorFinish = NULL;
static ExecutorEnd_hook_type prev_ExecutorEnd = NULL;
static ProcessUtility_hook_type prev_ProcessUtility = NULL;

/*---- GUC variables ----*/

typedef enum
{
	PGSP_TRACK_NONE,			/* track no plans */
	PGSP_TRACK_TOP,				/* only plans for top level statements */
	PGSP_TRACK_ALL,				/* all plans, including for nested statements */
}			PGSPTrackLevel;

static const struct config_enum_entry track_options[] =
{
	{"none", PGSP_TRACK_NONE, false},
	{"top", PGSP_TRACK_TOP, false},
	{"all", PGSP_TRACK_ALL, false},
	{NULL, 0, false}
};

typedef enum
{
	PGSP_COMPRESS_NONE,			/* no compression */
	PGSP_COMPRESS_ZSTD,			/* zstd compression */
}			PGSPCompression;

static const struct config_enum_entry compress_options[] =
{
	{"none", PGSP_COMPRESS_NONE, false},
	{"zstd", PGSP_COMPRESS_ZSTD, false},
	{NULL, 0, false}
};

static int	pgsp_max = 5000;	/* max # plans to track */
static int	pgsp_max_size = 8192;	/* max size of plan text to track (in
									 * bytes) */
static int	pgsp_track = PGSP_TRACK_TOP;	/* tracking level */
static int	pgsp_compress = PGSP_COMPRESS_NONE; /* compression type */
static int	pgsp_max_plan_memory = 16384;	/* plan text DSA size limit, in KB */
static bool pgsp_plan_advice = false;	/* collect pg_plan_advice strings */

/*
 * Module shared state, provisioned via the DSM registry (PG18+) or the shmem
 * hooks (PG<18).
 */
typedef struct pgspSharedState
{
	pg_atomic_uint64 nentries;

	/* Serialize garbage collection: only one backend runs it at a time. */
	pg_atomic_flag gc_active;

#if PG_VERSION_NUM < 190000
	/*
	 * Backing store for GetNamedDSA() shim on < PG19. Protected by AddinShmemInitLock.
	 */
	dsa_handle	plan_dsa_handle;	/* DSA_HANDLE_INVALID until created */
	int			plan_dsa_tranche;	/* 0 until allocated */
#endif
}			pgspSharedState;

static pgspSharedState *pgsp_shared = NULL;

/* Separate DSA holding plan texts, size-limited by pg_stat_plans.max_plan_memory */
static dsa_area *pgsp_plan_dsa = NULL;

#define pgsp_enabled(level) \
	(!IsParallelWorker() && \
	(pgsp_track == PGSP_TRACK_ALL || \
	(pgsp_track == PGSP_TRACK_TOP && (level) == 0)))

#define USAGE_INCREASE			0.5 /* increase by this each time we report
									 * stats */
#define USAGE_DECREASE_FACTOR	(0.99)	/* decreased every
										 * pgsp_dealloc_entries */
#define USAGE_DEALLOC_PERCENT	5	/* free this % of entries at once */

/*---- Function declarations ----*/

PG_FUNCTION_INFO_V1(pg_stat_plans_reset);
PG_FUNCTION_INFO_V1(pg_stat_plans_2_0);
PG_FUNCTION_INFO_V1(pg_stat_plans_2_0_get_activity);

/* Structures for statistics of plans */
typedef struct PgStatShared_PlanInfo
{
	/* key elements that identify a plan (together with the dboid) */
	uint64		planid;
	uint64		queryid;
	Oid			userid;			/* userid is tracked to allow users to see
								 * their own query plans */
	bool		toplevel;		/* query executed at top level */

	dsa_pointer plan_text;		/* pointer to DSA memory containing plan text,
								 * InvalidDsaPointer if the plan-text DSA was full
								 * when this entry was created */
	size_t		plan_text_size; /* size of stored plan text */
	bool		plan_text_truncated;	/* whether stored plan text was
										 * truncated due to text size limit */
	int			plan_text_compression;	/* plan compression used */
	int			plan_encoding;	/* plan text encoding */

	dsa_pointer plan_advice;	/* pointer into the plan-text DSA holding the
								 * pg_plan_advice string (Postgres 19+),
								 * InvalidDsaPointer if absent or the DSA was full */
	size_t		plan_advice_size;	/* size of stored advice (incl. NUL) */
}			PgStatShared_PlanInfo;

typedef struct PgStat_StatPlanEntry
{
	PgStat_Counter exec_count;
	double		exec_time;
	double		usage;			/* Usage factor of the entry, used to
								 * prioritize which plans to age out */

	/* Only used in shared structure, not in local pending stats */
	PgStatShared_PlanInfo info;
}			PgStat_StatPlanEntry;

typedef struct PgStatShared_Plan
{
	PgStatShared_Common header;
	PgStat_StatPlanEntry stats;
}			PgStatShared_Plan;

static bool plan_stats_flush_cb(PgStat_EntryRef *entry_ref, bool nowait);
static uint64 pgsp_calculate_plan_id(PlannedStmt *result);
static void pgsp_attach_shmem(void);

static const PgStat_KindInfo plan_stats = {
	.name = "plan_stats",
	.fixed_amount = false,

#if PG_VERSION_NUM >= 180000
	/*
	 * We currently don't write to a file since plan texts would get lost (and
	 * just the stats on their own aren't that useful)
	 */
	.write_to_file = false,
#endif

	/*
	 * Plan statistics are available system-wide to simplify monitoring
	 * scripts
	 */
	.accessed_across_databases = true,

	.shared_size = sizeof(PgStatShared_Plan),
	.shared_data_off = offsetof(PgStatShared_Plan, stats),
	.shared_data_len = sizeof(((PgStatShared_Plan *) 0)->stats),
	.pending_size = sizeof(PgStat_StatPlanEntry),
	.flush_pending_cb = plan_stats_flush_cb,
};

/*
 * Compute stats entry idx from query ID and plan ID with an 8-byte hash.
 *
 * Whilst we could theorically just use the plan ID here, we intentionally
 * add the query ID into the mix to ease interpreting the data in combination
 * with pg_stat_statements.
 */
#define PGSTAT_PLAN_IDX(query_id, plan_id, user_id, toplevel) hash_combine64(toplevel, hash_combine64(query_id, hash_combine64(plan_id, user_id)))

/*
 * Kind ID reserved for statistics of plans.
 */
#define PGSTAT_KIND_PLANS	27

/*
 * Callback for stats handling
 */
static bool
plan_stats_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
{
	PgStat_StatPlanEntry *localent;
	PgStatShared_Plan *shfuncent;

	localent = (PgStat_StatPlanEntry *) entry_ref->pending;
	shfuncent = (PgStatShared_Plan *) entry_ref->shared_stats;

	if (!pgstat_custom_lock_entry(entry_ref, nowait))
		return false;

	shfuncent->stats.exec_count += localent->exec_count;
	shfuncent->stats.exec_time += localent->exec_time;
	shfuncent->stats.usage += localent->usage;

	pgstat_custom_unlock_entry(entry_ref);

	return true;
}

static char *
pgsp_explain_plan(QueryDesc *queryDesc)
{
	ExplainState *es;
	StringInfo	es_str;

	es = NewExplainState();
	es_str = es->str;

	/*
	 * We turn off COSTS since identical planids may have very different
	 * costs, and it could be misleading to only show the first recorded
	 * plan's costs.
	 */
	es->costs = false;
	es->format = EXPLAIN_FORMAT_TEXT;

	ExplainBeginOutput(es);
	ExplainPrintPlan(es, queryDesc);
	ExplainEndOutput(es);

	/* Ignore trailing newline emitted by ExplainPrintPlan */
	if (es_str->len > 0)
		es_str->data[es_str->len - 1] = '\0';

	return es_str->data;
}

/* Copied from pgstat_shmem.c (not exported, and we can't use alternate APIs) */
static void
pgstat_free_entry(PgStatShared_HashEntry *shent, dshash_seq_status *hstat)
{
	dsa_pointer pdsa;
	PgStat_Kind kind = shent->key.kind;

	/*
	 * Fetch dsa pointer before deleting entry - that way we can free the
	 * memory after releasing the lock.
	 */
	pdsa = shent->body;

	if (!hstat)
		dshash_delete_entry(pgStatCustomLocal.shared_hash, shent);
	else
		dshash_delete_current(hstat);

	dsa_free(pgStatCustomLocal.dsa, pdsa);

#if PG_VERSION_NUM >= 190000
	/* Decrement entry count, if required. */
	if (pgstat_get_kind_info(kind)->track_entry_count)
		pg_atomic_sub_fetch_u64(&pgStatCustomLocal.shmem->entry_counts[kind - 1], 1);
#endif
}

typedef struct PlanDeallocEntry
{
	PgStat_HashKey key;
	double		usage;
}			PlanDeallocEntry;

/*
 * list sort comparator for sorting into decreasing usage order
 */
static int
entry_cmp_lru(const union ListCell *lhs, const union ListCell *rhs)
{
	double		l_usage = ((PlanDeallocEntry *) lfirst(lhs))->usage;
	double		r_usage = ((PlanDeallocEntry *) lfirst(rhs))->usage;

	if (l_usage > r_usage)
		return -1;
	else if (l_usage < r_usage)
		return +1;
	else
		return 0;
}

/* Deallocate n (USAGE_DEALLOC_PERCENT) entries based on their usage value */
static void
pgsp_dealloc_entries(void)
{
	dshash_seq_status hstat;
	PgStatShared_HashEntry *p;
	List	   *entries = NIL;
	ListCell   *lc;
	int			nvictims;

	/* dshash entry is not modified, take shared lock */
	dshash_seq_init(&hstat, pgStatCustomLocal.shared_hash, false);
	while ((p = dshash_seq_next(&hstat)) != NULL)
	{
		PgStatShared_Common *header;
		PgStat_StatPlanEntry *statent;
		PlanDeallocEntry *entry;

		if (p->dropped || p->key.kind != PGSTAT_KIND_PLANS)
			continue;

		header = dsa_get_address(pgStatCustomLocal.dsa, p->body);

		if (!LWLockConditionalAcquire(&header->lock, LW_EXCLUSIVE))
			continue;

		statent = (PgStat_StatPlanEntry *) pgstat_custom_get_entry_data(PGSTAT_KIND_PLANS, header);
		statent->usage *= USAGE_DECREASE_FACTOR;

		entry = palloc(sizeof(PlanDeallocEntry));
		entry->key = p->key;
		entry->usage = statent->usage;

		LWLockRelease(&header->lock);

		entries = lappend(entries, entry);
	}
	dshash_seq_term(&hstat);

	/* Sort by usage ascending (lowest used entries are last) */
	list_sort(entries, entry_cmp_lru);

	/* At a minimum, deallocate 10 entries to make it worth our while */
	nvictims = Max(10, list_length(entries) * USAGE_DEALLOC_PERCENT / 100);
	nvictims = Min(nvictims, list_length(entries));

	/* Actually drop the entries */
	for_each_from(lc, entries, list_length(entries) - nvictims)
	{
		PlanDeallocEntry *entry = lfirst(lc);
		PgStatShared_HashEntry *p;

		/*
		 * Free the plan text (and any advice string) before dropping the entry.
		 * If we left it for pgsp_dealloc_release() to free later, a concurrent
		 * resurrection of the entry by pgstat_reinit_entry() could cause a leak,
		 * since reinit does a memset and through that zeroes our DSA pointers.
		 */
		p = dshash_find(pgStatCustomLocal.shared_hash, &entry->key, true);
		if (p != NULL)
		{
			PgStatShared_Common *header =
				dsa_get_address(pgStatCustomLocal.dsa, p->body);
			PgStat_StatPlanEntry *statent =
				(PgStat_StatPlanEntry *) pgstat_custom_get_entry_data(PGSTAT_KIND_PLANS, header);

			LWLockAcquire(&header->lock, LW_EXCLUSIVE);
			if (DsaPointerIsValid(statent->info.plan_text))
			{
				dsa_free(pgsp_plan_dsa, statent->info.plan_text);
				statent->info.plan_text = InvalidDsaPointer;
			}
			if (DsaPointerIsValid(statent->info.plan_advice))
			{
				dsa_free(pgsp_plan_dsa, statent->info.plan_advice);
				statent->info.plan_advice = InvalidDsaPointer;
			}
			LWLockRelease(&header->lock);

			dshash_release_lock(pgStatCustomLocal.shared_hash, p);
		}

#if PG_VERSION_NUM >= 180000
		pgstat_custom_drop_entry(entry->key.kind, entry->key.dboid, entry->key.objid);
#else
		pgstat_custom_drop_entry(entry->key.kind, entry->key.dboid, entry->key.objoid);
#endif
	}

	/* Clean up our working memory immediately */
	foreach(lc, entries)
	{
		PlanDeallocEntry *entry = lfirst(lc);

		pfree(entry);
	}
	pfree(entries);
}

/*
 * Release (refcount to 0 + drop) entries after ensuring query text got freed.
 *
 * This is needed in addition to the entry walk (which also frees query text)
 * because certain operations (reset function, DROP DATABASE) cause us to have
 * entries that are dropped but whose query text has not been freed.
 */
static void
pgsp_dealloc_release(void)
{
	dshash_seq_status hstat;
	PgStatShared_HashEntry *p;
	List	   *victims = NIL;
	ListCell   *lc;

	/*
	 * Phase 1: find reclaimable entries under a *shared* scan.
	 *
	 * An entry is reclaimable once it is dropped and its refcount has fallen to
	 * 1 -- the single extra reference pg_stat_plans took at creation. A dropped
	 * entry can gain no new references. If a backend still holds a reference we
	 * skip it now and reclaim it on a later pass, once
	 * pgstat_request_entry_refs_gc() has nudged that backend to let go.
	 */
	dshash_seq_init(&hstat, pgStatCustomLocal.shared_hash, false);
	while ((p = dshash_seq_next(&hstat)) != NULL)
	{
		PgStat_HashKey *key;

		if (!p->dropped || p->key.kind != PGSTAT_KIND_PLANS)
			continue;
		if (pg_atomic_read_u32(&p->refcount) != 1)
			continue;

		key = palloc(sizeof(PgStat_HashKey));
		*key = p->key;
		victims = lappend(victims, key);
	}
	dshash_seq_term(&hstat);

	/*
	 * Phase 2: reclaim each victim under a brief exclusive lock on just its own
	 * partition. We re-check the entry under that lock, because between the two
	 * phases it may have been resurrected (reinitialized) by a new query reusing
	 * the same key.
	 */
	foreach(lc, victims)
	{
		PgStat_HashKey *key = (PgStat_HashKey *) lfirst(lc);
		PgStatShared_Common *header;
		PgStat_StatPlanEntry *statent;

		p = dshash_find(pgStatCustomLocal.shared_hash, key, true);
		if (p == NULL)
			continue;			/* already reclaimed by another backend */

		if (!p->dropped || p->key.kind != PGSTAT_KIND_PLANS ||
			pg_atomic_read_u32(&p->refcount) != 1)
		{
			/* referenced again or resurrected since phase 1 -- leave it */
			dshash_release_lock(pgStatCustomLocal.shared_hash, p);
			continue;
		}

		header = dsa_get_address(pgStatCustomLocal.dsa, p->body);
		statent = (PgStat_StatPlanEntry *) pgstat_custom_get_entry_data(PGSTAT_KIND_PLANS, header);

		/*
		 * Free the plan text (and any advice string), if still stored, from the
		 * DSA. We don't need a header lock here since we hold the partition lock
		 * exclusively and with refcount = 1 no backend holds a direct reference.
		 */
		if (DsaPointerIsValid(statent->info.plan_text))
		{
			dsa_free(pgsp_plan_dsa, statent->info.plan_text);
			statent->info.plan_text = InvalidDsaPointer;
		}
		if (DsaPointerIsValid(statent->info.plan_advice))
		{
			dsa_free(pgsp_plan_dsa, statent->info.plan_advice);
			statent->info.plan_advice = InvalidDsaPointer;
		}

		/* Drop our extra reference (1 -> 0) and free the entry. */
		pg_atomic_fetch_sub_u32(&p->refcount, 1);

		pgstat_free_entry(p, NULL);

		/* Work is done */
		pg_atomic_fetch_sub_u64(&pgsp_shared->nentries, 1);
	}

	list_free_deep(victims);

	/* Encourage other backends to clean up dropped entry refs */
	pgstat_custom_request_entry_refs_gc();
}

static void
pgsp_dealloc(void)
{
	/*
	 * Decide whether to run a cycle from the live-entry counter instead of
	 * scanning the whole hash to count. The counter is maintained on entry
	 * creation and in pgstat_gc_plan_memory(), and the caller invokes us with
	 * no entry lock held, so this is cheap and lock-order safe.
	 */
	if (pg_atomic_read_u64(&pgsp_shared->nentries) <= (uint64) pgsp_max)
		return;

	/*
	 * Serialize GC cycles. Multiple backends can be over the limit at once, as
	 * we trigger GC on every over-limit execution); if they all deallocated
	 * concurrently they would pick overlapping victims and the second drop of
	 * a victim would raise "trying to drop stats entry already dropped".
	 */
	if (!pg_atomic_test_set_flag(&pgsp_shared->gc_active))
		return;

	PG_TRY();
	{
		/* re-check now that we hold the GC slot */
		if (pg_atomic_read_u64(&pgsp_shared->nentries) > (uint64) pgsp_max)
		{
			pgsp_dealloc_entries();
			pgsp_dealloc_release();
		}
	}
	PG_FINALLY();
	{
		pg_atomic_clear_flag(&pgsp_shared->gc_active);
	}
	PG_END_TRY();
}

static char *
pgsp_compress_plan_text_zstd(char *plan, size_t plan_len, size_t *stored_plan_size, bool *stored_plan_truncated)
{
#ifdef USE_ZSTD
	ZSTD_CCtx  *cctx = ZSTD_createCCtx();
	ZSTD_outBuffer outBuf = {0};
	int			level = ZSTD_CLEVEL_DEFAULT;
	size_t		compressStep = 100; /* Compress in 100 byte input increments */
	size_t		yet_to_flush;

	size_t		ret = ZSTD_CCtx_setParameter(cctx, ZSTD_c_compressionLevel,
											 level);

	if (ZSTD_isError(ret))
	{
		elog(INFO, "could not set zstd compression level to %d: %s",
			 level, ZSTD_getErrorName(ret));
		return NULL;
	}

	/* Allocate the output buffer to its largest allowed extent, since we */
	/* re-allocate anyway when copying into DSA */
	outBuf.dst = palloc0(sizeof(char) * pgsp_max_size);
	outBuf.size = pgsp_max_size;
	outBuf.pos = 0;

	/*
	 * Compress data in steps, stopping early if we're likely to exceed our
	 * limit
	 */
	for (int pos = 0; pos < plan_len; pos += compressStep)
	{
		ZSTD_inBuffer inBuf = {plan + pos, Min(compressStep, plan_len - pos), 0};

		do
		{
			size_t		max_needed = ZSTD_compressBound(inBuf.size - inBuf.pos);

			if (outBuf.size - outBuf.pos < max_needed)
			{
				elog(INFO, "FIXME: unexpected short bound");
				return NULL;
			}

			yet_to_flush = ZSTD_compressStream2(cctx,
												&outBuf,
												&inBuf, ZSTD_e_continue);

			if (ZSTD_isError(yet_to_flush))
			{
				elog(INFO, "could not compress data: %s",
					 ZSTD_getErrorName(yet_to_flush));
				return NULL;
			}
		} while (yet_to_flush > 0);

		/*
		 * Stop early if output pos is close to max (to leave room for closing
		 * current frame)
		 */
		if ((outBuf.pos + compressStep) > outBuf.size * 0.95)
		{
			*stored_plan_truncated = true;
			break;
		}
	}

	/* Close current frame */
	do
	{
		ZSTD_inBuffer inBuf = {NULL, 0, 0};
		size_t		max_needed = ZSTD_compressBound(0);

		if (outBuf.size - outBuf.pos < max_needed)
		{
			elog(INFO, "FIXME: repalloc?");
			return NULL;
		}

		yet_to_flush = ZSTD_compressStream2(cctx,
											&outBuf,
											&inBuf, ZSTD_e_end);

		if (ZSTD_isError(yet_to_flush))
		{
			elog(INFO, "could not compress data: %s",
				 ZSTD_getErrorName(yet_to_flush));
			return NULL;
		}
	} while (yet_to_flush > 0);

	*stored_plan_size = outBuf.pos;

	return outBuf.dst;
#else
	elog(INFO, "zstd compression not enabled in build, falling back to no compression");
	return NULL;
#endif
}

static char *
pgsp_decompress_plan_text_zstd(char *stored_plan, size_t stored_plan_size)
{
#ifdef USE_ZSTD
	ZSTD_DCtx  *const dctx = ZSTD_createDCtx();
	ZSTD_inBuffer inBuf = {stored_plan, stored_plan_size, 0};
	ZSTD_outBuffer outBuf = {0};
	size_t		ret;

	outBuf.dst = palloc(sizeof(char) * pgsp_max_size * 20); /* TODO: Don't hardcode
															 * this */
	outBuf.size = pgsp_max_size * 20;

	ret = ZSTD_decompressStream(dctx, &outBuf, &inBuf);
	if (ZSTD_isError(ret))
		elog(ERROR, "could not decompress data: %s",
			 ZSTD_getErrorName(ret));
	if (ret != 0)
		elog(ERROR, "EOF before end of stream: %zu", ret);

	ZSTD_freeDCtx(dctx);

	((char *) outBuf.dst)[outBuf.pos] = '\0';

	return outBuf.dst;
#else
	elog(ERROR, "zstd compression not enabled in build (BUG: how did we get a plan that was marked as such?)");
	return NULL;
#endif
}

static char *
pgsp_compress_plan_text(char *plan, size_t *stored_plan_size, bool *stored_plan_truncated, int *stored_plan_compression)
{
	size_t		plan_len = strlen(plan);

	if (pgsp_compress == PGSP_COMPRESS_ZSTD)
	{
		char	   *stored_plan = pgsp_compress_plan_text_zstd(plan, plan_len, stored_plan_size, stored_plan_truncated);

		if (stored_plan != NULL)
		{
			*stored_plan_compression = PGSP_COMPRESS_ZSTD;
			return stored_plan;
		}
	}

	/* Default (no compression) */
	*stored_plan_size = Min(plan_len, pgsp_max_size);
	*stored_plan_truncated = plan_len > pgsp_max_size;
	*stored_plan_compression = PGSP_COMPRESS_NONE;

	return plan;
}

static char *
pgsp_decompress_plan_text(char *stored_plan, size_t stored_plan_size, int stored_plan_compression)
{
	switch (stored_plan_compression)
	{
		case PGSP_COMPRESS_ZSTD:
			{
				return pgsp_decompress_plan_text_zstd(stored_plan, stored_plan_size);
			}
		case PGSP_COMPRESS_NONE:
			{
				char	   *plan = palloc(sizeof(char) * stored_plan_size + 1);

				memcpy(plan, stored_plan, stored_plan_size);
				plan[stored_plan_size] = '\0';
				return plan;
			}
	}

	pg_unreachable();
}

#if PG_VERSION_NUM >= 190000

/*
 * Signature of pg_plan_advice_request_advice_generation(), resolved at runtime
 * so that we don't have a hard link-time dependency on pg_plan_advice.
 */
typedef void (*pgsp_request_advice_generation_fn) (bool activate);

static pgsp_request_advice_generation_fn pgsp_advice_fn = NULL;
static bool pgsp_advice_fn_resolved = false;

/*
 * Resolve pg_plan_advice_request_advice_generation() if pg_plan_advice is
 * loaded. We only attempt the lookup when pg_plan_advice is listed in
 * shared_preload_libraries, since load_external_function() would otherwise
 * raise an error when the library is not present. Returns NULL when advice
 * generation cannot be triggered.
 */
static pgsp_request_advice_generation_fn
pgsp_get_advice_fn(void)
{
	if (!pgsp_advice_fn_resolved)
	{
		const char *preload;

		pgsp_advice_fn_resolved = true;

		preload = GetConfigOption("shared_preload_libraries", true, false);
		if (preload != NULL && strstr(preload, "pg_plan_advice") != NULL)
			pgsp_advice_fn = (pgsp_request_advice_generation_fn)
				load_external_function("$libdir/pg_plan_advice",
									   "pg_plan_advice_request_advice_generation",
									   false, NULL);
	}

	return pgsp_advice_fn;
}

/*
 * Extract the generated pg_plan_advice string from a planned statement, if
 * pg_plan_advice stashed one in PlannedStmt.extension_state. Returns NULL when
 * no advice is present.
 */
static char *
pgsp_get_plan_advice(PlannedStmt *pstmt)
{
	ListCell   *lc;
	List	   *pgpa_list = NIL;

	foreach(lc, pstmt->extension_state)
	{
		DefElem    *item = (DefElem *) lfirst(lc);

		if (strcmp(item->defname, "pg_plan_advice") == 0)
		{
			pgpa_list = (List *) item->arg;
			break;
		}
	}

	foreach(lc, pgpa_list)
	{
		DefElem    *item = (DefElem *) lfirst(lc);

		if (strcmp(item->defname, "advice_string") == 0)
			return strVal(item->arg);
	}

	return NULL;
}

#endif							/* PG_VERSION_NUM >= 190000 */

static void
pgstat_report_plan_stats(QueryDesc *queryDesc,
						 PgStat_Counter exec_count,
						 double exec_time)
{
	PgStat_EntryRef *entry_ref;
	PgStatShared_Plan *shstatent;
	PgStat_StatPlanEntry *pending;
	bool		newly_created;
	uint64		queryId = queryDesc->plannedstmt->queryId;
	uint64		planId;
	Oid			userid = GetUserId();
	bool		toplevel = (nesting_level == 0);

#if PG_VERSION_NUM >= 180000
	planId = queryDesc->plannedstmt->planId;
#else
	planId = pgsp_calculate_plan_id(queryDesc->plannedstmt);
#endif

	/* Make sure our shared state is attached before we touch the counter. */
	pgsp_attach_shmem();

	/*
	 * Run garbage collection if needed before creating a new entry. Doing it
	 * now, rather than inline while holding the new entry's header lock, means
	 * GC's hash-wide scans never run under a held entry lock, and the entry
	 * created here cannot be chosen as a victim because it does not exist yet.
	 */
	if (pgsp_shared != NULL &&
		pg_atomic_read_u64(&pgsp_shared->nentries) > (uint64) pgsp_max)
		pgsp_dealloc();

	entry_ref = pgstat_custom_prep_pending_entry(PGSTAT_KIND_PLANS, MyDatabaseId,
										  PGSTAT_PLAN_IDX(queryId, planId, userid, toplevel), &newly_created);

	shstatent = (PgStatShared_Plan *) entry_ref->shared_stats;
	pending = (PgStat_StatPlanEntry *) entry_ref->pending;

	if (newly_created)
	{
		char	   *plan = pgsp_explain_plan(queryDesc);
		size_t		stored_plan_size = 0;
		bool		stored_plan_truncated = false;
		int			stored_plan_compression = 0;
		char	   *stored_plan = pgsp_compress_plan_text(plan, &stored_plan_size, &stored_plan_truncated, &stored_plan_compression);

		(void) pgstat_custom_lock_entry(entry_ref, false);

		shstatent->stats.info.planid = planId;
		shstatent->stats.info.queryid = queryId;
		shstatent->stats.info.userid = userid;
		shstatent->stats.info.toplevel = toplevel;

		/*
		 * Store the plan text in our separate, size-limited DSA. If that DSA is
		 * full, keep the entry but without text (DSA_ALLOC_NO_OOM returns
		 * InvalidDsaPointer rather than erroring) -- the entry's counters are
		 * still useful and the SRF renders a NULL plan.
		 */
		shstatent->stats.info.plan_text =
			dsa_allocate_extended(pgsp_plan_dsa, stored_plan_size, DSA_ALLOC_NO_OOM);
		if (DsaPointerIsValid(shstatent->stats.info.plan_text))
		{
			memcpy(dsa_get_address(pgsp_plan_dsa, shstatent->stats.info.plan_text), stored_plan, stored_plan_size);
			shstatent->stats.info.plan_text_size = stored_plan_size;
			shstatent->stats.info.plan_text_truncated = stored_plan_truncated;
			shstatent->stats.info.plan_text_compression = stored_plan_compression;
		}
		else
		{
			shstatent->stats.info.plan_text_size = 0;
			shstatent->stats.info.plan_text_truncated = false;
			shstatent->stats.info.plan_text_compression = 0;
		}
		shstatent->stats.info.plan_encoding = GetDatabaseEncoding();
		shstatent->stats.info.plan_advice = InvalidDsaPointer;
		shstatent->stats.info.plan_advice_size = 0;

		/*
		 * Take an extra reference so pgstat can't free the entry (and orphan
		 * its plan text) until pgsp_dealloc_release() reclaims it. We do this
		 * for every new entry, even ones without text, so that we have an
		 * accurate total entry count, and gc_plan_memory is solely responsible
		 * to ensure final cleanup.
		 *
		 * The exception to this are resurrected entries, which can occur when
		 * an entry was marked as dropped, but not actually removed yet. If
		 * a dropped entry is attempted to be recreated, the old entry gets
		 * reset, but the reference counts are kept. We can detect this through
		 * the generation counter, and pretend no new entry was created.
		 */
		if (pg_atomic_read_u32(&entry_ref->shared_entry->generation) == 0)
		{
			pg_atomic_fetch_add_u32(&entry_ref->shared_entry->refcount, 1);
			pg_atomic_fetch_add_u64(&pgsp_shared->nentries, 1);
		}

		pgstat_custom_unlock_entry(entry_ref);

		pfree(plan);
	}

#if PG_VERSION_NUM >= 190000
	/*
	 * Record the pg_plan_advice string for this plan, if we don't have it yet.
	 * The advice is generated during planning (when we requested it via
	 * pgsp_planner) and stashed in PlannedStmt.extension_state.
	 */
	if (pgsp_plan_advice &&
		!DsaPointerIsValid(shstatent->stats.info.plan_advice))
	{
		char	   *advice = pgsp_get_plan_advice(queryDesc->plannedstmt);

		if (advice != NULL && advice[0] != '\0')
		{
			size_t		advice_size = strlen(advice) + 1;

			(void) pgstat_custom_lock_entry(entry_ref, false);

			/* re-check under lock, another backend may have stored it */
			if (!DsaPointerIsValid(shstatent->stats.info.plan_advice))
			{
				/*
				 * Store the advice in the same size-limited DSA as the plan
				 * text. If that DSA is full, keep the entry but without advice
				 * (DSA_ALLOC_NO_OOM returns InvalidDsaPointer rather than
				 * erroring).
				 */
				dsa_pointer advice_dp =
					dsa_allocate_extended(pgsp_plan_dsa, advice_size, DSA_ALLOC_NO_OOM);

				if (DsaPointerIsValid(advice_dp))
				{
					memcpy(dsa_get_address(pgsp_plan_dsa, advice_dp), advice, advice_size);
					shstatent->stats.info.plan_advice = advice_dp;
					shstatent->stats.info.plan_advice_size = advice_size;
				}
			}

			pgstat_custom_unlock_entry(entry_ref);
		}
	}
#endif

	pending->exec_count += exec_count;
	pending->exec_time += exec_time;
	pending->usage += USAGE_INCREASE;
}

static void
pgsp_plan_id_walker(JumbleState *jstate, Plan *plan)
{
	ListCell   *l;

	if (plan == NULL)
		return;

	/*
	 * Plan-type-specific walks
	 */
	switch (nodeTag(plan))
	{
		case T_SubqueryScan:
			{
				SubqueryScan *sscan = (SubqueryScan *) plan;

				pgsp_plan_id_walker(jstate, sscan->subplan);
			}
			break;
		case T_CustomScan:
			{
				CustomScan *cscan = (CustomScan *) plan;

				foreach(l, cscan->custom_plans)
				{
					pgsp_plan_id_walker(jstate, (Plan *) lfirst(l));
				}
			}
			break;
		case T_Append:
			{
				Append	   *aplan = (Append *) plan;

				foreach(l, aplan->appendplans)
				{
					pgsp_plan_id_walker(jstate, (Plan *) lfirst(l));
				}
			}
			break;
		case T_MergeAppend:
			{
				MergeAppend *mplan = (MergeAppend *) plan;

				foreach(l, mplan->mergeplans)
				{
					pgsp_plan_id_walker(jstate, (Plan *) lfirst(l));
				}
			}
			break;
		case T_BitmapAnd:
			{
				BitmapAnd  *splan = (BitmapAnd *) plan;

				foreach(l, splan->bitmapplans)
				{
					pgsp_plan_id_walker(jstate, (Plan *) lfirst(l));
				}
			}
			break;
		case T_BitmapOr:
			{
				BitmapOr   *splan = (BitmapOr *) plan;

				foreach(l, splan->bitmapplans)
				{
					pgsp_plan_id_walker(jstate, (Plan *) lfirst(l));
				}
			}
			break;
		default:
			{
				/* do nothing */
			}
	}

	pgsp_plan_id_walker(jstate, plan->lefttree);
	pgsp_plan_id_walker(jstate, plan->righttree);

	JumbleNode(jstate, (Node *) plan);
}

static uint64
pgsp_calculate_plan_id(PlannedStmt *result)
{
	JumbleState *jstate = InitJumble();
	ListCell   *lc;
	uint64 planId;

	pgsp_plan_id_walker(jstate, result->planTree);
	foreach(lc, result->subplans)
	{
		Plan	   *subplan = (Plan *) lfirst(lc);

		pgsp_plan_id_walker(jstate, subplan);
	}

	JumbleRangeTable(jstate, result->rtable);
	planId = HashJumbleState(jstate);
	pfree(jstate);
	return planId;
}

/*
 * Planner hook: forward to regular planner, but increase plan count and
 * record query plan if needed.
 */
static PlannedStmt *
#if PG_VERSION_NUM >= 190000
pgsp_planner(Query *parse,
			 const char *query_string,
			 int cursorOptions,
			 ParamListInfo boundParams,
			 ExplainState *es)
#else
pgsp_planner(Query *parse,
			 const char *query_string,
			 int cursorOptions,
			 ParamListInfo boundParams)
#endif
{
	PlannedStmt *result;
#if PG_VERSION_NUM >= 190000
	bool		capture_advice = false;
	pgsp_request_advice_generation_fn advice_fn = NULL;

	/*
	 * Decide whether to ask pg_plan_advice to generate a plan advice string
	 * for this query. We only do this for queries we would actually track. The
	 * advice ends up in PlannedStmt.extension_state and is recorded in
	 * pgstat_report_plan_stats.
	 *
	 * Note: nesting_level has not been incremented yet, so pgsp_enabled()
	 * reflects this query, matching the check done after planning below.
	 */
	if (pgsp_plan_advice &&
		parse->queryId != UINT64CONST(0) &&
		pgsp_enabled(nesting_level))
	{
		advice_fn = pgsp_get_advice_fn();
		if (advice_fn != NULL)
			capture_advice = true;
	}

	if (capture_advice)
		advice_fn(true);
#endif

	/*
	 * Increment the nesting level, to ensure that functions evaluated during
	 * planning are not seen as top-level calls.
	 */
	nesting_level++;
	PG_TRY();
	{
#if PG_VERSION_NUM >= 190000
		if (prev_planner_hook)
			result = prev_planner_hook(parse, query_string, cursorOptions,
									   boundParams, es);
		else
			result = standard_planner(parse, query_string, cursorOptions,
									  boundParams, es);
#else
		if (prev_planner_hook)
			result = prev_planner_hook(parse, query_string, cursorOptions,
									   boundParams);
		else
			result = standard_planner(parse, query_string, cursorOptions,
									  boundParams);
#endif
	}
	PG_FINALLY();
	{
		nesting_level--;
#if PG_VERSION_NUM >= 190000
		if (capture_advice)
			advice_fn(false);
#endif
	}
	PG_END_TRY();

#if PG_VERSION_NUM >= 180000
	if (pgsp_enabled(nesting_level))
		result->planId = pgsp_calculate_plan_id(result);
#endif

	return result;
}


/*
 * Post-parse-analysis hook: Reset query ID to support EXECUTE statements
 */
static void
#if PG_VERSION_NUM >= 190000
pgsp_post_parse_analyze(ParseState *pstate, Query *query, const JumbleState *jstate)
#else
pgsp_post_parse_analyze(ParseState *pstate, Query *query, JumbleState *jstate)
#endif
{
	if (prev_post_parse_analyze_hook)
		prev_post_parse_analyze_hook(pstate, query, jstate);

	/* Safety check... */
	if (!pgsp_enabled(nesting_level))
		return;

	/*
	 * If it's EXECUTE, clear the queryId. This matches pg_stat_statements and
	 * is necessary to make pg_stat_plans_activity work correctly for EXECUTE
	 * statements when pg_stat_statements is not loaded.
	 */
	if (query->utilityStmt && IsA(query->utilityStmt, ExecuteStmt))
		query->queryId = UINT64CONST(0);
}

/*
 * ExecutorStart hook: start up tracking if needed
 */
static void
pgsp_ExecutorStart(QueryDesc *queryDesc, int eflags)
{
	uint64		queryId = queryDesc->plannedstmt->queryId;
#if PG_VERSION_NUM >= 180000
	uint64		planId = queryDesc->plannedstmt->planId;
#endif

#if PG_VERSION_NUM >= 190000
	/* Set up to track total elapsed time in ExecutorRun. */
	if (queryId != UINT64CONST(0) &&
		planId != UINT64CONST(0) &&
		pgsp_enabled(nesting_level))
		queryDesc->query_instr_options |= INSTRUMENT_TIMER;
#endif

	if (prev_ExecutorStart)
		prev_ExecutorStart(queryDesc, eflags);
	else
		standard_ExecutorStart(queryDesc, eflags);

	if (queryId != UINT64CONST(0) &&
#if PG_VERSION_NUM >= 180000
		planId != UINT64CONST(0) &&
#endif
		pgsp_enabled(nesting_level))
	{
#if PG_VERSION_NUM >= 180000
		/*
		 * Record initial entry now, so plan text is available for currently
		 * running queries
		 */
		pgstat_report_plan_stats(queryDesc,
								 0, /* executions are counted in
									 * pgsp_ExecutorEnd */
								 0.0);
#endif

#if PG_VERSION_NUM < 190000
		/*
		 * Set up to track total elapsed time in ExecutorRun.  Make sure the
		 * space is allocated in the per-query context so it will go away at
		 * ExecutorEnd.
		 */
		if (queryDesc->totaltime == NULL)
		{
			MemoryContext oldcxt;

			oldcxt = MemoryContextSwitchTo(queryDesc->estate->es_query_cxt);
			queryDesc->totaltime = InstrAlloc(1, INSTRUMENT_ALL, false);
			MemoryContextSwitchTo(oldcxt);
		}
#endif
	}
}

/*
 * ExecutorRun hook: all we need do is track nesting depth
 */
static void
#if PG_VERSION_NUM >= 180000
pgsp_ExecutorRun(QueryDesc *queryDesc, ScanDirection direction, uint64 count)
#else
pgsp_ExecutorRun(QueryDesc *queryDesc, ScanDirection direction, uint64 count, bool execute_once)
#endif
{
	nesting_level++;
	PG_TRY();
	{
#if PG_VERSION_NUM >= 180000
		if (prev_ExecutorRun)
			prev_ExecutorRun(queryDesc, direction, count);
		else
			standard_ExecutorRun(queryDesc, direction, count);
#else
		if (prev_ExecutorRun)
			prev_ExecutorRun(queryDesc, direction, count, execute_once);
		else
			standard_ExecutorRun(queryDesc, direction, count, execute_once);
#endif
	}
	PG_FINALLY();
	{
		nesting_level--;
	}
	PG_END_TRY();
}

/*
 * ExecutorFinish hook: all we need do is track nesting depth
 */
static void
pgsp_ExecutorFinish(QueryDesc *queryDesc)
{
	nesting_level++;
	PG_TRY();
	{
		if (prev_ExecutorFinish)
			prev_ExecutorFinish(queryDesc);
		else
			standard_ExecutorFinish(queryDesc);
	}
	PG_FINALLY();
	{
		nesting_level--;
	}
	PG_END_TRY();
}

/*
 * ExecutorEnd hook: store results if needed
 */
static void
pgsp_ExecutorEnd(QueryDesc *queryDesc)
{
	uint64		queryId = queryDesc->plannedstmt->queryId;
#if PG_VERSION_NUM >= 180000
	uint64		planId = queryDesc->plannedstmt->planId;
#endif

	if (queryId != UINT64CONST(0) &&
#if PG_VERSION_NUM >= 180000
		planId != UINT64CONST(0) &&
#endif
#if PG_VERSION_NUM >= 190000
		queryDesc->query_instr && pgsp_enabled(nesting_level))
#else
		queryDesc->totaltime && pgsp_enabled(nesting_level))
#endif
	{
#if PG_VERSION_NUM < 190000
		/*
		 * Make sure stats accumulation is done.  (Note: it's okay if several
		 * levels of hook all do this.)
		 */
		InstrEndLoop(queryDesc->totaltime);

		pgstat_report_plan_stats(queryDesc,
								 1,
								 queryDesc->totaltime->total * 1000.0 /* convert to msec */ );
#else
		pgstat_report_plan_stats(queryDesc,
								 1,
								 INSTR_TIME_GET_MILLISEC(queryDesc->query_instr->total));
#endif

#if PG_VERSION_NUM < 180000
		/* TODO: Is there a better way to do this on < PG18? */
		pgstat_custom_report_stat(true);
#endif
	}

	if (prev_ExecutorEnd)
		prev_ExecutorEnd(queryDesc);
	else
		standard_ExecutorEnd(queryDesc);
}

/*
 * ProcessUtility hook
 */
static void
pgsp_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
					bool readOnlyTree,
					ProcessUtilityContext context,
					ParamListInfo params, QueryEnvironment *queryEnv,
					DestReceiver *dest, QueryCompletion *qc)
{
	Node	   *parsetree = pstmt->utilityStmt;
	bool		enabled = pgsp_enabled(nesting_level);

	/*
	 * Even though we're not tracking plans for utility statements themselves,
	 * we must still increment the nesting level, to ensure that functions
	 * evaluated within it are not seen as top-level calls.  But don't do so
	 * for EXECUTE; that way, when control reaches pgsp_planner or
	 * pgsp_ExecutorStart, we will treat the costs as top-level if
	 * appropriate.  Likewise, don't bump for PREPARE, so that parse analysis
	 * will treat the statement as top-level if appropriate.
	 *
	 * To be absolutely certain we don't mess up the nesting level, evaluate
	 * the bump_level condition just once.
	 */
	bool		bump_level =
		!IsA(parsetree, ExecuteStmt) &&
		!IsA(parsetree, PrepareStmt);

	/*
	 * Force utility statements to get queryId zero. This matches what
	 * pg_stat_statements does in its ProcessUtility hook, and is necessary to
	 * allow joining on the queryId for pg_stat_plans_activity when
	 * pg_stat_statements is not enabled. If pg_stat_statements is loaded this
	 * is effectively a no-op, since either before our call of the hook, or
	 * afterwards, pgss_ProcessUtility will do the same thing.
	 */
	if (enabled)
		pstmt->queryId = UINT64CONST(0);

	if (bump_level)
		nesting_level++;
	PG_TRY();
	{
		if (prev_ProcessUtility)
			prev_ProcessUtility(pstmt, queryString, readOnlyTree,
								context, params, queryEnv,
								dest, qc);
		else
			standard_ProcessUtility(pstmt, queryString, readOnlyTree,
									context, params, queryEnv,
									dest, qc);
	}
	PG_FINALLY();
	{
		if (bump_level)
			nesting_level--;
	}
	PG_END_TRY();
}

#if PG_VERSION_NUM < 190000
/*
 * Minimal back-port of PG19's GetNamedDSA() for our single plan-text DSA. The
 * handle and LWLock tranche id live in our own shared state (PG19 keeps them in
 * the core DSM registry); one backend creates the area, the rest attach. *found
 * is false only on first creation, so the caller applies the size limit once.
 */
static dsa_area *
pgsp_named_dsa(bool *found)
{
	dsa_area   *area;
	MemoryContext oldcontext;

	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);

	if (pgsp_shared->plan_dsa_tranche == 0)
		pgsp_shared->plan_dsa_tranche = LWLockNewTrancheId();

	/* The tranche id is shared, but its name table is backend-local. */
	LWLockRegisterTranche(pgsp_shared->plan_dsa_tranche, "pg_stat_plans_plan");

	/*
	 * dsa_create()/dsa_attach() palloc the backend-local dsa_area in the current
	 * memory context, and dsa_pin_mapping() registers a cleanup tied to it.
	 * Switch to TopMemoryContext so they live for the whole backend, matching
	 * what GetNamedDSA() does on PG19.
	 */
	oldcontext = MemoryContextSwitchTo(TopMemoryContext);

	if (pgsp_shared->plan_dsa_handle == DSA_HANDLE_INVALID)
	{
		*found = false;
		area = dsa_create(pgsp_shared->plan_dsa_tranche);
		dsa_pin(area);
		dsa_pin_mapping(area);
		pgsp_shared->plan_dsa_handle = dsa_get_handle(area);
	}
	else
	{
		*found = true;
		area = dsa_attach(pgsp_shared->plan_dsa_handle);
		dsa_pin_mapping(area);
	}

	MemoryContextSwitchTo(oldcontext);

	LWLockRelease(AddinShmemInitLock);
	return area;
}
#endif

/*
 * Attach to (creating on first use) the separate, size-limited DSA that holds
 * plan texts. On PG19 this is core's GetNamedDSA(); on older versions it is the
 * shim above. The size limit is a soft cap that can be raised/lowered at
 * runtime, so the GUC is PGC_SIGHUP (assign hook below).
 */
static void
pgsp_attach_plan_dsa(void)
{
	bool		found;

	if (pgsp_plan_dsa != NULL)
		return;

#if PG_VERSION_NUM >= 190000
	pgsp_plan_dsa = GetNamedDSA("pg_stat_plans_plan", &found);
#else
	pgsp_plan_dsa = pgsp_named_dsa(&found);
#endif

	if (!found)
		dsa_set_size_limit(pgsp_plan_dsa, (size_t) pgsp_max_plan_memory * 1024);
}

/* SIGHUP assign hook: move the live plan-text memory limit. */
static void
pgsp_assign_max_plan_memory(int newval, void *extra)
{
	if (pgsp_plan_dsa != NULL)
		dsa_set_size_limit(pgsp_plan_dsa, (size_t) newval * 1024);
}

#if PG_VERSION_NUM >= 180000

/*
 * One-time initializer for our shared state, run by the DSM registry under its
 * own lock the first time any backend attaches.
 */
#if PG_VERSION_NUM >= 190000
static void
pgsp_init_shmem(void *ptr, void *arg)
#else
static void
pgsp_init_shmem(void *ptr)
#endif
{
	pgspSharedState *state = (pgspSharedState *) ptr;

	pg_atomic_init_u64(&state->nentries, 0);
	pg_atomic_init_flag(&state->gc_active);
#if PG_VERSION_NUM < 190000
	state->plan_dsa_handle = DSA_HANDLE_INVALID;
	state->plan_dsa_tranche = 0;
#endif
}

/*
 * Lazily attach to our shared state via the DSM registry. Called at the top of
 * the entry points that touch pgsp_shared / pgsp_plan_dsa. Mirrors
 * pg_stat_statements: no shmem_request/shmem_startup hooks are needed -- the
 * registry creates the segment on first use and hands every backend the same
 * pointer.
 */
static void
pgsp_attach_shmem(void)
{
	if (pgsp_shared == NULL)
	{
		bool		found;

#if PG_VERSION_NUM >= 190000
		pgsp_shared = GetNamedDSMSegment("pg_stat_plans",
										 sizeof(pgspSharedState),
										 pgsp_init_shmem, &found, NULL);
#else
		pgsp_shared = GetNamedDSMSegment("pg_stat_plans",
										 sizeof(pgspSharedState),
										 pgsp_init_shmem, &found);
#endif
	}

	pgsp_attach_plan_dsa();
}

#else							/* PG_VERSION_NUM < 180000 */

/*
 * Before PG18 the custom cumulative stats machinery is backported in the
 * compat shim, which still needs the classic shmem hooks; our shared state
 * rides along in the same startup callback rather than using the DSM registry.
 */
static shmem_request_hook_type prev_shmem_request_hook = NULL;
static shmem_startup_hook_type prev_shmem_startup_hook = NULL;

static void
pgsp_shmem_request(void)
{
	if (prev_shmem_request_hook)
		prev_shmem_request_hook();

	RequestAddinShmemSpace(MAXALIGN(sizeof(pgspSharedState)));
}

static void
pgsp_shmem_startup(void)
{
	bool		found;

	if (prev_shmem_startup_hook)
		prev_shmem_startup_hook();

	StatsCustomShmemInit();

	/* Reset in case this is a restart within the postmaster */
	pgsp_shared = NULL;
	pgsp_plan_dsa = NULL;

	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
	pgsp_shared = ShmemInitStruct("pg_stat_plans",
								  sizeof(pgspSharedState), &found);
	if (!found)
	{
		pg_atomic_init_u64(&pgsp_shared->nentries, 0);
		pg_atomic_init_flag(&pgsp_shared->gc_active);
		pgsp_shared->plan_dsa_handle = DSA_HANDLE_INVALID;
		pgsp_shared->plan_dsa_tranche = 0;
	}
	LWLockRelease(AddinShmemInitLock);
}

/* The startup hook set pgsp_shared; attach the (lazily created) plan DSA. */
static void
pgsp_attach_shmem(void)
{
	pgsp_attach_plan_dsa();
}

static ClientAuthentication_hook_type prev_ClientAuthentication_hook = NULL;

static void
pgsp_ClientAuthentication_hook(Port *port, int status)
{
	if (prev_ClientAuthentication_hook)
		prev_ClientAuthentication_hook(port, status);

	pgstat_custom_initialize();
}
#endif

/*
 * Module load callback
 */
void
_PG_init(void)
{
	/*
	 * In order to register for shared memory stats, we have to be loaded via
	 * shared_preload_libraries.  If not, fall out without hooking into any of
	 * the main system.  (We don't throw error here because it seems useful to
	 * allow the pg_stat_plans functions to be created even when the module
	 * isn't active.  The functions must protect themselves against being
	 * called then, however.)
	 */
	if (!process_shared_preload_libraries_in_progress)
		return;

	/*
	 * Inform the postmaster that we want to enable query_id calculation if
	 * compute_query_id is set to auto.
	 */
	EnableQueryId();

	/*
	 * Define (or redefine) custom GUC variables.
	 */
	DefineCustomIntVariable("pg_stat_plans.max",
							"Sets the maximum number of plans tracked by pg_stat_plans in shared memory.",
							NULL,
							&pgsp_max,
							5000,
							100,
							INT_MAX / 2,
							PGC_SIGHUP,
							0,
							NULL,
							NULL,
							NULL);

	DefineCustomIntVariable("pg_stat_plans.max_size",
							"Sets the maximum size of plan texts (in bytes) tracked by pg_stat_plans in shared memory.",
							NULL,
							&pgsp_max_size,
							8192,
							100,
							1048576,	/* 1MB hard limit */
							PGC_SUSET,
							0,
							NULL,
							NULL,
							NULL);

	DefineCustomIntVariable("pg_stat_plans.max_plan_memory",
							"Sets the memory limit for plan text storage.",
							NULL,
							&pgsp_max_plan_memory,
							16384,
							256,
							MAX_KILOBYTES,
							PGC_SIGHUP,
							GUC_UNIT_KB,
							NULL,
							pgsp_assign_max_plan_memory,
							NULL);

	DefineCustomEnumVariable("pg_stat_plans.track",
							 "Selects which plans are tracked by pg_stat_plans.",
							 NULL,
							 &pgsp_track,
							 PGSP_TRACK_TOP,
							 track_options,
							 PGC_SUSET,
							 0,
							 NULL,
							 NULL,
							 NULL);

	DefineCustomEnumVariable("pg_stat_plans.compress",
							 "Select compression used by pg_stat_plans.",
							 NULL,
							 &pgsp_compress,
							 PGSP_COMPRESS_NONE,
							 compress_options,
							 PGC_SUSET,
							 0,
							 NULL,
							 NULL,
							 NULL);

	DefineCustomBoolVariable("pg_stat_plans.plan_advice",
							 "Collect pg_plan_advice strings for tracked plans.",
							 "Only effective on Postgres 19 or newer with pg_plan_advice loaded.",
							 &pgsp_plan_advice,
							 false,
							 PGC_SUSET,
							 0,
							 NULL,
							 NULL,
							 NULL);

	MarkGUCPrefixReserved("pg_stat_plans");

	/*
	 * Install hooks.
	 */
	prev_post_parse_analyze_hook = post_parse_analyze_hook;
	post_parse_analyze_hook = pgsp_post_parse_analyze;
	prev_planner_hook = planner_hook;
	planner_hook = pgsp_planner;
	prev_ExecutorStart = ExecutorStart_hook;
	ExecutorStart_hook = pgsp_ExecutorStart;
	prev_ExecutorRun = ExecutorRun_hook;
	ExecutorRun_hook = pgsp_ExecutorRun;
	prev_ExecutorFinish = ExecutorFinish_hook;
	ExecutorFinish_hook = pgsp_ExecutorFinish;
	prev_ExecutorEnd = ExecutorEnd_hook;
	ExecutorEnd_hook = pgsp_ExecutorEnd;
	prev_ProcessUtility = ProcessUtility_hook;
	ProcessUtility_hook = pgsp_ProcessUtility;

#if PG_VERSION_NUM < 180000
	prev_shmem_request_hook = shmem_request_hook;
	shmem_request_hook = pgsp_shmem_request;
	prev_shmem_startup_hook = shmem_startup_hook;
	shmem_startup_hook = pgsp_shmem_startup;
	prev_ClientAuthentication_hook = ClientAuthentication_hook;
	ClientAuthentication_hook = pgsp_ClientAuthentication_hook;
#endif

	pgstat_custom_register_kind(PGSTAT_KIND_PLANS, &plan_stats);
}

static bool
match_plans_entries(PgStatShared_HashEntry *entry, Datum match_data)
{
	return entry->key.kind == PGSTAT_KIND_PLANS;
}

/*
 * Reset statement statistics.
 */
Datum
pg_stat_plans_reset(PG_FUNCTION_ARGS)
{
	/* stats kind must be registered already */
	if (!pgstat_custom_get_kind_info(PGSTAT_KIND_PLANS))
		ereport(ERROR,
				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
				 errmsg("pg_stat_plans must be loaded via \"shared_preload_libraries\"")));

	/* Make sure our shared state is attached before gc updates the counter. */
	pgsp_attach_shmem();

	pgstat_custom_drop_matching_entries(match_plans_entries, 0);

	/* Free plan text memory and allow cleanup of dropped entries */
	pgsp_dealloc_release();

	PG_RETURN_VOID();
}

Datum
pg_stat_plans_2_0(PG_FUNCTION_ARGS)
{
#define PG_STAT_PLANS_COLS 9
	bool		showplan = PG_GETARG_BOOL(0);
	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
	Oid			userid = GetUserId();
	bool		is_allowed_role = false;

	dshash_seq_status hstat;
	PgStatShared_HashEntry *p;

	/*
	 * Superusers or roles with the privileges of pg_read_all_stats members
	 * are allowed
	 */
	is_allowed_role = has_privs_of_role(userid, ROLE_PG_READ_ALL_STATS);

	/* stats kind must be registered already */
	if (!pgstat_custom_get_kind_info(PGSTAT_KIND_PLANS))
		ereport(ERROR,
				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
				 errmsg("pg_stat_plans must be loaded via \"shared_preload_libraries\"")));

	/* Ensure the plan-text DSA is attached before we read plan texts. */
	pgsp_attach_shmem();

	InitMaterializedSRF(fcinfo, 0);

	/* dshash entry is not modified, take shared lock */
	dshash_seq_init(&hstat, pgStatCustomLocal.shared_hash, false);
	while ((p = dshash_seq_next(&hstat)) != NULL)
	{
		PgStat_StatPlanEntry *statent;
		Datum		values[PG_STAT_PLANS_COLS];
		bool		nulls[PG_STAT_PLANS_COLS];
		int			i = 0;

		if (p->dropped || p->key.kind != PGSTAT_KIND_PLANS)
			continue;

		memset(values, 0, sizeof(values));
		memset(nulls, 0, sizeof(nulls));

		statent = pgstat_custom_get_entry_data(p->key.kind, dsa_get_address(pgStatCustomLocal.dsa, p->body));

		values[i++] = ObjectIdGetDatum(statent->info.userid);
		values[i++] = ObjectIdGetDatum(p->key.dboid);
		values[i++] = BoolGetDatum(statent->info.toplevel);
		if (is_allowed_role || statent->info.userid == userid)
		{
			int64		queryid = statent->info.queryid;
			int64		planid = statent->info.planid;

			values[i++] = Int64GetDatumFast(queryid);
			values[i++] = Int64GetDatumFast(planid);
		}
		else
		{
			nulls[i++] = true;
			nulls[i++] = true;
		}
		values[i++] = Int64GetDatumFast(statent->exec_count);
		values[i++] = Float8GetDatumFast(statent->exec_time);

		if (showplan && (is_allowed_role || statent->info.userid == userid))
		{
			char	   *pstr = DsaPointerIsValid(statent->info.plan_text) ? dsa_get_address(pgsp_plan_dsa, statent->info.plan_text) : NULL;

			if (pstr)
			{
				char	   *dec_pstr = pgsp_decompress_plan_text(pstr, statent->info.plan_text_size, statent->info.plan_text_compression);
				char	   *enc = pg_any_to_server(dec_pstr, strlen(dec_pstr), statent->info.plan_encoding);

				values[i++] = CStringGetTextDatum(enc);

				if (enc != dec_pstr)
					pfree(enc);
				pfree(dec_pstr);
			}
			else
			{
				nulls[i++] = true;
			}
		}
		else if (showplan)
		{
			values[i++] = CStringGetTextDatum("<insufficient privilege>");
		}
		else
		{
			nulls[i++] = true;
		}

		/* plan_advice (Postgres 19+; always NULL otherwise) */
#if PG_VERSION_NUM >= 190000
		if (showplan && (is_allowed_role || statent->info.userid == userid) &&
			DsaPointerIsValid(statent->info.plan_advice))
		{
			char	   *astr = dsa_get_address(pgsp_plan_dsa, statent->info.plan_advice);

			values[i++] = CStringGetTextDatum(astr);
		}
		else
		{
			nulls[i++] = true;
		}
#else
		nulls[i++] = true;
#endif

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

	return (Datum) 0;
}

#define HAS_PGSTAT_PERMISSIONS(role)	 (has_privs_of_role(GetUserId(), ROLE_PG_READ_ALL_STATS) || has_privs_of_role(GetUserId(), role))

/*
 * Returns plans of active PG backends.
 */
Datum
pg_stat_plans_2_0_get_activity(PG_FUNCTION_ARGS)
{
#if PG_VERSION_NUM >= 180000
#define PG_STAT_PLANS_GET_ACTIVITY_COLS	5
	int			num_backends = pgstat_fetch_stat_numbackends();
	int			curr_backend;
	int			pid = PG_ARGISNULL(0) ? -1 : PG_GETARG_INT32(0);
	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;

	InitMaterializedSRF(fcinfo, 0);

	/* 1-based index */
	for (curr_backend = 1; curr_backend <= num_backends; curr_backend++)
	{
		/* for each row */
		Datum		values[PG_STAT_PLANS_GET_ACTIVITY_COLS] = {0};
		bool		nulls[PG_STAT_PLANS_GET_ACTIVITY_COLS] = {0};
		LocalPgBackendStatus *local_beentry;
		PgBackendStatus *beentry;

		/* Get the next one in the list */
		local_beentry = pgstat_get_local_beentry_by_index(curr_backend);
		beentry = &local_beentry->backendStatus;

		/* If looking for specific PID, ignore all the others */
		if (pid != -1 && beentry->st_procpid != pid)
			continue;

		/* Values available to all callers */
		if (beentry->st_databaseid != InvalidOid)
			values[0] = ObjectIdGetDatum(beentry->st_databaseid);
		else
			nulls[0] = true;

		values[1] = Int32GetDatum(beentry->st_procpid);

		if (beentry->st_userid != InvalidOid)
			values[2] = ObjectIdGetDatum(beentry->st_userid);
		else
			nulls[2] = true;

		/* Values only available to role member or pg_read_all_stats */
		if (HAS_PGSTAT_PERMISSIONS(beentry->st_userid))
		{
			if (beentry->st_query_id == 0)
				nulls[3] = true;
			else
				values[3] = UInt64GetDatum(beentry->st_query_id);

			if (beentry->st_plan_id == 0)
				nulls[4] = true;
			else
				values[4] = UInt64GetDatum(beentry->st_plan_id);
		}
		else
		{
			/* No permissions to view data about this session */
			nulls[3] = true;
			nulls[4] = true;
		}

		tuplestore_putvalues(rsinfo->setResult, rsinfo->setDesc, values, nulls);

		/* If only a single backend was requested, and we found it, break. */
		if (pid != -1)
			break;
	}
#else
	elog(ERROR, "Not implemented, use of pg_stat_plans_get_activity requires Postgres 18+");
#endif

	return (Datum) 0;
}
