Huh, maybe that's the better way to do it?

Author: lfittl

src/backend/access/transam/xact.c

@@ -40,6 +40,7 @@
 #include "commands/tablecmds.h"
 #include "commands/trigger.h"
 #include "common/pg_prng.h"
+#include "executor/instrument.h"
 #include "executor/spi.h"
 #include "libpq/be-fsstubs.h"
 #include "libpq/pqsignal.h"
@@ -2937,6 +2938,7 @@ AbortTransaction(void)
 	 * do abort processing
 	 */
 	AfterTriggerEndXact(false); /* 'false' means it's abort */
+	AtAbort_Instrumentation();
 	AtAbort_Portals();
 	smgrDoPendingSyncs(false, is_parallel_worker);
 	AtEOXact_LargeObject(false);
@@ -5338,6 +5340,7 @@ AbortSubTransaction(void)
 	if (s->curTransactionOwner)
 	{
 		AfterTriggerEndSubXact(false);
+		AtSubAbort_Instrumentation(s->nestingLevel);
 		AtSubAbort_Portals(s->subTransactionId,
 						   s->parent->subTransactionId,
 						   s->curTransactionOwner,

src/backend/executor/instrument.c

@@ -15,14 +15,17 @@
 
 #include <unistd.h>
 
+#include "access/xact.h"
 #include "executor/instrument.h"
 #include "utils/memutils.h"
-#include "utils/resowner.h"
 
 WalUsage	pgWalUsage;
 Instrumentation instr_top;
 InstrStackState instr_stack = {0, 0, NULL, &instr_top};
 
+/* List of active QueryInstrumentation entries, for abort cleanup */
+static dlist_head active_query_instrumentations = DLIST_STATIC_INIT(active_query_instrumentations);
+
 void
 InstrStackGrow(void)
 {
@@ -98,7 +101,7 @@ InstrStop(Instrumentation *instr)
  * Stops instrumentation, finalizes the stack entry and accumulates to its parent.
  *
  * Note that this intentionally allows passing a stack that is not the current
- * top, as can happen with PG_FINALLY, or resource owners, which don't have a
+ * top, as can happen with PG_FINALLY or abort cleanup, which don't have a
  * guaranteed cleanup order.
  *
  * We are careful here to achieve two goals:
@@ -140,34 +143,16 @@ InstrStopFinalize(Instrumentation *instr)
 /* Query instrumentation handling */
 
 /*
- * Use ResourceOwner mechanism to correctly reset instr_stack on abort.
+ * Finalize a single QueryInstrumentation entry during abort.
+ *
+ * Accumulates buffer/WAL data from unfinalized children and triggers,
+ * resets the instrumentation stack, and removes the entry from the
+ * active list. We do NOT pfree any memory here since portal/executor
+ * context cleanup will handle that shortly after.
  */
-static void ResOwnerReleaseInstrumentation(Datum res);
-static const ResourceOwnerDesc instrumentation_resowner_desc =
-{
-	.name = "instrumentation",
-	.release_phase = RESOURCE_RELEASE_AFTER_LOCKS,
-	.release_priority = RELEASE_PRIO_INSTRUMENTATION,
-	.ReleaseResource = ResOwnerReleaseInstrumentation,
-	.DebugPrint = NULL,			/* default message is fine */
-};
-
-static inline void
-ResourceOwnerRememberInstrumentation(ResourceOwner owner, QueryInstrumentation *qinstr)
-{
-	ResourceOwnerRemember(owner, PointerGetDatum(qinstr), &instrumentation_resowner_desc);
-}
-
-static inline void
-ResourceOwnerForgetInstrumentation(ResourceOwner owner, QueryInstrumentation *qinstr)
-{
-	ResourceOwnerForget(owner, PointerGetDatum(qinstr), &instrumentation_resowner_desc);
-}
-
 static void
-ResOwnerReleaseInstrumentation(Datum res)
+InstrAbortQueryInstrumentation(QueryInstrumentation *qinstr)
 {
-	QueryInstrumentation *qinstr = (QueryInstrumentation *) DatumGetPointer(res);
 	dlist_mutable_iter iter;
 
 	/* Accumulate data from all unfinalized child node entries. */
@@ -176,12 +161,6 @@ ResOwnerReleaseInstrumentation(Datum res)
 		NodeInstrumentation *child = dlist_container(NodeInstrumentation, unfinalized_node, iter.cur);
 
 		InstrAccumStack(&qinstr->instr, &child->instr);
-
-		/*
-		 * Free NodeInstrumentation now, since InstrFinalizeNode won't be
-		 * called
-		 */
-		pfree(child);
 	}
 
 	/* Accumulate data from any active trigger instrumentation entries. */
@@ -190,37 +169,66 @@ ResOwnerReleaseInstrumentation(Datum res)
 		TriggerInstrumentation *tginstr = dlist_container(TriggerInstrumentation, unfinalized_trigger, iter.cur);
 
 		InstrAccumStack(&qinstr->instr, &tginstr->instr);
-
-		/*
-		 * We can't pfree tginstr here, since its part of a bigger allocation - we'll instead
-		 * let transaction end deal with clean up, and instead just remove the list entry.
-		 */
-		dlist_delete(&tginstr->unfinalized_trigger);
 	}
 
 	/* Ensure the stack is reset as expected, and we accumulate to the parent */
 	InstrStopFinalize(&qinstr->instr);
+}
+
+/*
+ * Transaction abort cleanup for instrumentation.
+ *
+ * Called from AbortTransaction before AtAbort_Portals, so that we can
+ * recover buffer/WAL statistics from active instrumentation entries before
+ * the executor memory contexts are destroyed.
+ */
+void
+AtAbort_Instrumentation(void)
+{
+	dlist_mutable_iter iter;
+
+	dlist_foreach_modify(iter, &active_query_instrumentations)
+	{
+		QueryInstrumentation *qinstr = dlist_container(QueryInstrumentation, active_node, iter.cur);
+
+		InstrAbortQueryInstrumentation(qinstr);
+	}
+
+	/* Reset the list — memory will be freed by context cleanup */
+	dlist_init(&active_query_instrumentations);
+}
+
+/*
+ * Subtransaction abort cleanup for instrumentation.
+ *
+ * Only finalizes entries created at or below the given nesting level.
+ */
+void
+AtSubAbort_Instrumentation(int nestLevel)
+{
+	dlist_mutable_iter iter;
+
+	dlist_foreach_modify(iter, &active_query_instrumentations)
+	{
+		QueryInstrumentation *qinstr = dlist_container(QueryInstrumentation, active_node, iter.cur);
+
+		if (qinstr->nestingLevel < nestLevel)
+			continue;
 
-	/* Free QueryInstrumentation now, since InstrStop won't be called */
-	pfree(qinstr);
+		InstrAbortQueryInstrumentation(qinstr);
+		dlist_delete(&qinstr->active_node);
+	}
 }
 
 QueryInstrumentation *
 InstrQueryAlloc(int instrument_options)
 {
 	QueryInstrumentation *instr;
 
-	/*
-	 * If needed, allocate in TopMemoryContext so that the Instrumentation
-	 * survives transaction abort — ResourceOwner release needs to access
-	 * it.
-	 */
-	if ((instrument_options & (INSTRUMENT_BUFFERS | INSTRUMENT_WAL)) != 0)
-		instr = MemoryContextAllocZero(TopMemoryContext, sizeof(QueryInstrumentation));
-	else
-		instr = palloc0(sizeof(QueryInstrumentation));
+	instr = palloc0(sizeof(QueryInstrumentation));
 
 	InstrInitOptions(&instr->instr, instrument_options);
+	dlist_node_init(&instr->active_node);
 	dlist_init(&instr->unfinalized_children);
 	dlist_init(&instr->unfinalized_triggers);
 
@@ -232,58 +240,35 @@ InstrQueryStart(QueryInstrumentation *qinstr)
 {
 	InstrStart(&qinstr->instr);
 
-	if (qinstr->instr.need_stack)
+	if (qinstr->instr.need_stack && dlist_node_is_detached(&qinstr->active_node))
 	{
-		Assert(CurrentResourceOwner != NULL);
-		qinstr->owner = CurrentResourceOwner;
-
-		ResourceOwnerEnlarge(qinstr->owner);
-		ResourceOwnerRememberInstrumentation(qinstr->owner, qinstr);
+		qinstr->nestingLevel = GetCurrentTransactionNestLevel();
+		dlist_push_head(&active_query_instrumentations, &qinstr->active_node);
 	}
 }
 
 void
 InstrQueryStop(QueryInstrumentation *qinstr)
 {
 	InstrStop(&qinstr->instr);
-
-	if (qinstr->instr.need_stack)
-	{
-		Assert(qinstr->owner != NULL);
-		ResourceOwnerForgetInstrumentation(qinstr->owner, qinstr);
-		qinstr->owner = NULL;
-	}
 }
 
 QueryInstrumentation *
 InstrQueryStopFinalize(QueryInstrumentation *qinstr)
 {
-	QueryInstrumentation *copy;
-
 	InstrStopFinalize(&qinstr->instr);
 
-	if (!qinstr->instr.need_stack)
-		return qinstr;
-
-	Assert(qinstr->owner != NULL);
-	ResourceOwnerForgetInstrumentation(qinstr->owner, qinstr);
-	qinstr->owner = NULL;
+	if (qinstr->instr.need_stack && !dlist_node_is_detached(&qinstr->active_node))
+		dlist_delete_thoroughly(&qinstr->active_node);
 
-	/*
-	 * Copy to the current memory context so the caller doesn't need to
-	 * explicitly free the TopMemoryContext allocation.
-	 */
-	copy = palloc(sizeof(QueryInstrumentation));
-	memcpy(copy, qinstr, sizeof(QueryInstrumentation));
-	pfree(qinstr);
-	return copy;
+	return qinstr;
 }
 
 /*
  * Register a child NodeInstrumentation entry for abort processing.
  *
- * On abort, ResOwnerReleaseInstrumentation will walk the parent's list to
- * recover buffer/WAL data from entries that were never finalized, in order for
+ * On abort, AtAbort_Instrumentation will walk the parent's list to recover
+ * buffer/WAL data from entries that were never finalized, in order for
  * aggregate totals to be accurate despite the query erroring out.
  */
 void
@@ -340,17 +325,7 @@ InstrAllocNode(int instrument_options, bool async_mode)
 {
 	NodeInstrumentation *instr;
 
-	/*
-	 * If needed, allocate in a context that supports stack-based
-	 * instrumentation abort handling. We can utilize TopTransactionContext
-	 * instead of TopMemoryContext here because nodes don't get used for
-	 * utility commands that restart transactions, which would require a
-	 * context that survives longer (EXPLAIN ANALYZE is fine).
-	 */
-	if ((instrument_options & (INSTRUMENT_BUFFERS | INSTRUMENT_WAL)) != 0)
-		instr = MemoryContextAlloc(TopTransactionContext, sizeof(NodeInstrumentation));
-	else
-		instr = palloc(sizeof(NodeInstrumentation));
+	instr = palloc(sizeof(NodeInstrumentation));
 
 	InstrInitNode(instr, instrument_options);
 	instr->async_mode = async_mode;
@@ -422,26 +397,17 @@ InstrStopNode(NodeInstrumentation *instr, double nTuples)
 NodeInstrumentation *
 InstrFinalizeNode(NodeInstrumentation *instr, Instrumentation *parent)
 {
-	NodeInstrumentation *dst;
-
 	/* If we didn't use stack based instrumentation, nothing to be done */
 	if (!instr->instr.need_stack)
 		return instr;
 
-	/* Copy into per-query memory context */
-	dst = palloc(sizeof(NodeInstrumentation));
-	memcpy(dst, instr, sizeof(NodeInstrumentation));
-
 	/* Accumulate node's buffer/WAL usage to the parent */
-	InstrAccumStack(parent, &dst->instr);
+	InstrAccumStack(parent, &instr->instr);
 
-	/* Unregister from query's unfinalized list before freeing */
-	if (instr->instr.need_stack)
-		dlist_delete(&instr->unfinalized_node);
-
-	pfree(instr);
+	/* Unregister from query's unfinalized list */
+	dlist_delete(&instr->unfinalized_node);
 
-	return dst;
+	return instr;
 }
 
 /* Update tuple count */
@@ -630,18 +596,10 @@ TriggerInstrumentation *
 InstrAllocTrigger(int n, int instrument_options)
 {
 	TriggerInstrumentation *tginstr;
-
-	/*
-	 * If needed, allocate in TopTransactionContext so the memory survives
-	 * transaction abort — ResOwnerReleaseInstrumentation needs to access it.
-	 */
-	if ((instrument_options & (INSTRUMENT_BUFFERS | INSTRUMENT_WAL)) != 0)
-		tginstr = MemoryContextAllocZero(TopTransactionContext,
-										 n * sizeof(TriggerInstrumentation));
-	else
-		tginstr = palloc0(n * sizeof(TriggerInstrumentation));
 	int			i;
 
+	tginstr = palloc0(n * sizeof(TriggerInstrumentation));
+
 	for (i = 0; i < n; i++)
 		InstrInitOptions(&tginstr[i].instr, instrument_options);
 
@@ -657,7 +615,7 @@ InstrStartTrigger(QueryInstrumentation *qinstr, TriggerInstrumentation *tginstr)
 	 * On first call, register with the parent QueryInstrumentation for abort
 	 * recovery. The trigger stays on the list for the query lifetime -- on
 	 * normal completion ExecFinalizeTriggerInstrumentation handles it, on
-	 * abort ResOwnerReleaseInstrumentation does.
+	 * abort AtAbort_Instrumentation does.
 	 *
 	 * We detect first call by checking if the dlist_node is still in its
 	 * palloc0-zeroed state (prev == NULL).