1 files changed, 157 insertions, 0 deletions
diff --git a/src/backend/utils/adt/mcxtfuncs.c b/src/backend/utils/adt/mcxtfuncs.c
new file mode 100644
index 00000000000..50e1b07ff02
--- /dev/null
+++ b/src/backend/utils/adt/mcxtfuncs.c
@@ -0,0 +1,157 @@
+/*-------------------------------------------------------------------------
+ *
+ * mcxtfuncs.c
+ *	  Functions to show backend memory context.
+ *
+ * Portions Copyright (c) 1996-2020, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ *
+ * IDENTIFICATION
+ *	  src/backend/utils/adt/mcxtfuncs.c
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "funcapi.h"
+#include "miscadmin.h"
+#include "mb/pg_wchar.h"
+#include "utils/builtins.h"
+
+/* ----------
+ * The max bytes for showing identifiers of MemoryContext.
+ * ----------
+ */
+#define MEMORY_CONTEXT_IDENT_DISPLAY_SIZE	1024
+
+/*
+ * PutMemoryContextsStatsTupleStore
+ *		One recursion level for pg_get_backend_memory_contexts.
+ */
+static void
+PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
+								TupleDesc tupdesc, MemoryContext context,
+								const char *parent, int level)
+{
+#define PG_GET_BACKEND_MEMORY_CONTEXTS_COLS	9
+
+	Datum		values[PG_GET_BACKEND_MEMORY_CONTEXTS_COLS];
+	bool		nulls[PG_GET_BACKEND_MEMORY_CONTEXTS_COLS];
+	MemoryContextCounters stat;
+	MemoryContext child;
+	const char *name;
+	const char *ident;
+
+	AssertArg(MemoryContextIsValid(context));
+
+	name = context->name;
+	ident = context->ident;
+
+	/*
+	 * To be consistent with logging output, we label dynahash contexts
+	 * with just the hash table name as with MemoryContextStatsPrint().
+	 */
+	if (ident && strcmp(name, "dynahash") == 0)
+	{
+		name = ident;
+		ident = NULL;
+	}
+
+	/* Examine the context itself */
+	memset(&stat, 0, sizeof(stat));
+	(*context->methods->stats) (context, NULL, (void *) &level, &stat);
+
+	memset(values, 0, sizeof(values));
+	memset(nulls, 0, sizeof(nulls));
+
+	if (name)
+		values[0] = CStringGetTextDatum(name);
+	else
+		nulls[0] = true;
+
+	if (ident)
+	{
+		int		idlen = strlen(ident);
+		char		clipped_ident[MEMORY_CONTEXT_IDENT_DISPLAY_SIZE];
+
+		/*
+		 * Some identifiers such as SQL query string can be very long,
+		 * truncate oversize identifiers.
+		 */
+		if (idlen >= MEMORY_CONTEXT_IDENT_DISPLAY_SIZE)
+			idlen = pg_mbcliplen(ident, idlen, MEMORY_CONTEXT_IDENT_DISPLAY_SIZE - 1);
+
+		memcpy(clipped_ident, ident, idlen);
+		clipped_ident[idlen] = '\0';
+		values[1] = CStringGetTextDatum(clipped_ident);
+	}
+	else
+		nulls[1] = true;
+
+	if (parent)
+		values[2] = CStringGetTextDatum(parent);
+	else
+		nulls[2] = true;
+
+	values[3] = Int32GetDatum(level);
+	values[4] = Int64GetDatum(stat.totalspace);
+	values[5] = Int64GetDatum(stat.nblocks);
+	values[6] = Int64GetDatum(stat.freespace);
+	values[7] = Int64GetDatum(stat.freechunks);
+	values[8] = Int64GetDatum(stat.totalspace - stat.freespace);
+	tuplestore_putvalues(tupstore, tupdesc, values, nulls);
+
+	for (child = context->firstchild; child != NULL; child = child->nextchild)
+	{
+		PutMemoryContextsStatsTupleStore(tupstore, tupdesc,
+								  child, name, level + 1);
+	}
+}
+
+/*
+ * pg_get_backend_memory_contexts
+ *		SQL SRF showing backend memory context.
+ */
+Datum
+pg_get_backend_memory_contexts(PG_FUNCTION_ARGS)
+{
+	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
+	TupleDesc	tupdesc;
+	Tuplestorestate *tupstore;
+	MemoryContext per_query_ctx;
+	MemoryContext oldcontext;
+
+	/* check to see if caller supports us returning a tuplestore */
+	if (rsinfo == NULL || !IsA(rsinfo, ReturnSetInfo))
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("set-valued function called in context that cannot accept a set")));
+	if (!(rsinfo->allowedModes & SFRM_Materialize))
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("materialize mode required, but it is not allowed in this context")));
+
+	/* Build a tuple descriptor for our result type */
+	if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
+		elog(ERROR, "return type must be a row type");
+
+	per_query_ctx = rsinfo->econtext->ecxt_per_query_memory;
+	oldcontext = MemoryContextSwitchTo(per_query_ctx);
+
+	tupstore = tuplestore_begin_heap(true, false, work_mem);
+	rsinfo->returnMode = SFRM_Materialize;
+	rsinfo->setResult = tupstore;
+	rsinfo->setDesc = tupdesc;
+
+	MemoryContextSwitchTo(oldcontext);
+
+	PutMemoryContextsStatsTupleStore(tupstore, tupdesc,
+								TopMemoryContext, NULL, 0);
+
+	/* clean up and return the tuplestore */
+	tuplestore_donestoring(tupstore);
+
+	return (Datum) 0;
+}