/*-------------------------------------------------------------------------
 *
 * mcxtfuncs.c
 *   Functions to show backend memory context.
 *
 * Portions Copyright (c) 1996-2025, 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 "mb/pg_wchar.h"
#include "storage/proc.h"
#include "storage/procarray.h"
#include "utils/array.h"
#include "utils/builtins.h"
#include "utils/hsearch.h"

/* ----------
 * The max bytes for showing identifiers of MemoryContext.
 * ----------
 */
#define MEMORY_CONTEXT_IDENT_DISPLAY_SIZE 1024

/*
 * MemoryContextId
 *  Used for storage of transient identifiers for
 *  pg_get_backend_memory_contexts.
 */
typedef struct MemoryContextId
{
 MemoryContext context;
 int   context_id;
} MemoryContextId;

/*
 * int_list_to_array
 *  Convert an IntList to an array of INT4OIDs.
 */
static Datum
int_list_to_array(const List *list)
{
 Datum    *datum_array;
 int   length;
 ArrayType  *result_array;

 length = list_length(list);
 datum_array = (Datum *) palloc(length * sizeof(Datum));

 foreach_int(i, list)
  datum_array[foreach_current_index(i)] = Int32GetDatum(i);

 result_array = construct_array_builtin(datum_array, length, INT4OID);

 return PointerGetDatum(result_array);
}

/*
 * PutMemoryContextsStatsTupleStore
 *  Add details for the given MemoryContext to 'tupstore'.
 */
static void
PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
         TupleDesc tupdesc, MemoryContext context,
         HTAB *context_id_lookup)
{
#define PG_GET_BACKEND_MEMORY_CONTEXTS_COLS 10

 Datum  values[PG_GET_BACKEND_MEMORY_CONTEXTS_COLS];
 bool  nulls[PG_GET_BACKEND_MEMORY_CONTEXTS_COLS];
 MemoryContextCounters stat;
 List    *path = NIL;
 const char *name;
 const char *ident;
 const char *type;

 Assert(MemoryContextIsValid(context));

 /*
  * Figure out the transient context_id of this context and each of its
  * ancestors.
 */
 for (MemoryContext cur = context; cur != NULL; cur = cur->parent)
 {
  MemoryContextId *entry;
  bool  found;

  entry = hash_search(context_id_lookup, &cur, HASH_FIND, &found);

  if (!found)
   elog(ERROR, "hash table corrupted");
  path = lcons_int(entry->context_id, path);
 }

 /* Examine the context itself */
 memset(&stat, 0, sizeof(stat));
 (*context->methods->stats) (context, NULL, NULL, &stat, true);

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

 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;
 }

 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;

 switch (context->type)
 {
  case T_AllocSetContext:
   type = "AllocSet";
   break;
  case T_GenerationContext:
   type = "Generation";
   break;
  case T_SlabContext:
   type = "Slab";
   break;
  case T_BumpContext:
   type = "Bump";
   break;
  default:
   type = "???";
   break;
 }

 values[2] = CStringGetTextDatum(type);
 values[3] = Int32GetDatum(list_length(path)); /* level */
 values[4] = int_list_to_array(path);
 values[5] = Int64GetDatum(stat.totalspace);
 values[6] = Int64GetDatum(stat.nblocks);
 values[7] = Int64GetDatum(stat.freespace);
 values[8] = Int64GetDatum(stat.freechunks);
 values[9] = Int64GetDatum(stat.totalspace - stat.freespace);

 tuplestore_putvalues(tupstore, tupdesc, values, nulls);
 list_free(path);
}

/*
 * 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;
 int   context_id;
 List    *contexts;
 HASHCTL  ctl;
 HTAB    *context_id_lookup;

 ctl.keysize = sizeof(MemoryContext);
 ctl.entrysize = sizeof(MemoryContextId);
 ctl.hcxt = CurrentMemoryContext;

 context_id_lookup = hash_create("pg_get_backend_memory_contexts",
         256,
         &ctl,
         HASH_ELEM | HASH_BLOBS | HASH_CONTEXT);

 InitMaterializedSRF(fcinfo, 0);

 /*
  * Here we use a non-recursive algorithm to visit all MemoryContexts
  * starting with TopMemoryContext.  The reason we avoid using a recursive
  * algorithm is because we want to assign the context_id breadth-first.
  * I.e. all contexts at level 1 are assigned IDs before contexts at level
  * 2.  Because contexts closer to TopMemoryContext are less likely to
  * change, this makes the assigned context_id more stable.  Otherwise, if
  * the first child of TopMemoryContext obtained an additional grandchild,
  * the context_id for the second child of TopMemoryContext would change.
 */
 contexts = list_make1(TopMemoryContext);

 /* TopMemoryContext will always have a context_id of 1 */
 context_id = 1;

 foreach_ptr(MemoryContextData, cur, contexts)
 {
  MemoryContextId *entry;
  bool  found;

  /*
   * Record the context_id that we've assigned to each MemoryContext.
   * PutMemoryContextsStatsTupleStore needs this to populate the "path"
   * column with the parent context_ids.
 */
  entry = (MemoryContextId *) hash_search(context_id_lookup, &cur,
            HASH_ENTER, &found);
  entry->context_id = context_id++;
  Assert(!found);

  PutMemoryContextsStatsTupleStore(rsinfo->setResult,
           rsinfo->setDesc,
           cur,
           context_id_lookup);

  /*
   * Append all children onto the contexts list so they're processed by
   * subsequent iterations.
 */
  for (MemoryContext c = cur->firstchild; c != NULL; c = c->nextchild)
   contexts = lappend(contexts, c);
 }

 hash_destroy(context_id_lookup);

 return (Datum) 0;
}

/*
 * pg_log_backend_memory_contexts
 *  Signal a backend or an auxiliary process to log its memory contexts.
 *
 * By default, only superusers are allowed to signal to log the memory
 * contexts because allowing any users to issue this request at an unbounded
 * rate would cause lots of log messages and which can lead to denial of
 * service. Additional roles can be permitted with GRANT.
 *
 * On receipt of this signal, a backend or an auxiliary process sets the flag
 * in the signal handler, which causes the next CHECK_FOR_INTERRUPTS()
 * or process-specific interrupt handler to log the memory contexts.
 */
Datum
pg_log_backend_memory_contexts(PG_FUNCTION_ARGS)
{
 int   pid = PG_GETARG_INT32(0);
 PGPROC    *proc;
 ProcNumber procNumber = INVALID_PROC_NUMBER;

 /*
  * See if the process with given pid is a backend or an auxiliary process.
 */
 proc = BackendPidGetProc(pid);
 if (proc == NULL)
  proc = AuxiliaryPidGetProc(pid);

 /*
  * BackendPidGetProc() and AuxiliaryPidGetProc() return NULL if the pid
  * isn't valid; but by the time we reach kill(), a process for which we
  * get a valid proc here might have terminated on its own.  There's no way
  * to acquire a lock on an arbitrary process to prevent that. But since
  * this mechanism is usually used to debug a backend or an auxiliary
  * process running and consuming lots of memory, that it might end on its
  * own first and its memory contexts are not logged is not a problem.
 */
 if (proc == NULL)
 {
  /*
   * This is just a warning so a loop-through-resultset will not abort
   * if one backend terminated on its own during the run.
 */
  ereport(WARNING,
    (errmsg("PID %d is not a PostgreSQL server process", pid)));
  PG_RETURN_BOOL(false);
 }

 procNumber = GetNumberFromPGProc(proc);
 if (SendProcSignal(pid, PROCSIG_LOG_MEMORY_CONTEXT, procNumber) < 0)
 {
  /* Again, just a warning to allow loops */
  ereport(WARNING,
    (errmsg("could not send signal to process %d: %m", pid)));
  PG_RETURN_BOOL(false);
 }

 PG_RETURN_BOOL(true);
}