Invent a memory context reset/delete callback mechanism.

This allows cleanup actions to be registered to be called just before a
particular memory context's contents are flushed (either by deletion or
MemoryContextReset).  The patch in itself has no use-cases for this, but
several likely reasons for wanting this exist.

In passing, per discussion, rearrange some boolean fields in struct
MemoryContextData so as to avoid wasted padding space.  For safety,
this requires making allowInCritSection's existence unconditional;
but I think that's a better approach than what was there anyway.

Author: tglsfdc

src/backend/utils/mmgr/mcxt.c

@@ -54,6 +54,7 @@ MemoryContext CurTransactionContext = NULL;
 /* This is a transient link to the active portal's memory context: */
 MemoryContext PortalContext = NULL;
 
+static void MemoryContextCallResetCallbacks(MemoryContext context);
 static void MemoryContextStatsInternal(MemoryContext context, int level);
 
 /*
@@ -115,9 +116,8 @@ MemoryContextInit(void)
 	 * where retained memory in a context is *essential* --- we want to be
 	 * sure ErrorContext still has some memory even if we've run out
 	 * elsewhere! Also, allow allocations in ErrorContext within a critical
-	 * section. Otherwise a PANIC will cause an assertion failure in the
-	 * error reporting code, before printing out the real cause of the
-	 * failure.
+	 * section. Otherwise a PANIC will cause an assertion failure in the error
+	 * reporting code, before printing out the real cause of the failure.
 	 *
 	 * This should be the last step in this function, as elog.c assumes memory
 	 * management works once ErrorContext is non-null.
@@ -150,6 +150,7 @@ MemoryContextReset(MemoryContext context)
 	/* Nothing to do if no pallocs since startup or last reset */
 	if (!context->isReset)
 	{
+		MemoryContextCallResetCallbacks(context);
 		(*context->methods->reset) (context);
 		context->isReset = true;
 		VALGRIND_DESTROY_MEMPOOL(context);
@@ -195,6 +196,14 @@ MemoryContextDelete(MemoryContext context)
 
 	MemoryContextDeleteChildren(context);
 
+	/*
+	 * It's not entirely clear whether 'tis better to do this before or after
+	 * delinking the context; but an error in a callback will likely result in
+	 * leaking the whole context (if it's not a root context) if we do it
+	 * after, so let's do it before.
+	 */
+	MemoryContextCallResetCallbacks(context);
+
 	/*
 	 * We delink the context from its parent before deleting it, so that if
 	 * there's an error we won't have deleted/busted contexts still attached
@@ -242,6 +251,56 @@ MemoryContextResetAndDeleteChildren(MemoryContext context)
 	MemoryContextReset(context);
 }
 
+/*
+ * MemoryContextRegisterResetCallback
+ *		Register a function to be called before next context reset/delete.
+ *		Such callbacks will be called in reverse order of registration.
+ *
+ * The caller is responsible for allocating a MemoryContextCallback struct
+ * to hold the info about this callback request, and for filling in the
+ * "func" and "arg" fields in the struct to show what function to call with
+ * what argument.  Typically the callback struct should be allocated within
+ * the specified context, since that means it will automatically be freed
+ * when no longer needed.
+ *
+ * There is no API for deregistering a callback once registered.  If you
+ * want it to not do anything anymore, adjust the state pointed to by its
+ * "arg" to indicate that.
+ */
+void
+MemoryContextRegisterResetCallback(MemoryContext context,
+								   MemoryContextCallback *cb)
+{
+	AssertArg(MemoryContextIsValid(context));
+
+	/* Push onto head so this will be called before older registrants. */
+	cb->next = context->reset_cbs;
+	context->reset_cbs = cb;
+	/* Mark the context as non-reset (it probably is already). */
+	context->isReset = false;
+}
+
+/*
+ * MemoryContextCallResetCallbacks
+ *		Internal function to call all registered callbacks for context.
+ */
+static void
+MemoryContextCallResetCallbacks(MemoryContext context)
+{
+	MemoryContextCallback *cb;
+
+	/*
+	 * We pop each callback from the list before calling.  That way, if an
+	 * error occurs inside the callback, we won't try to call it a second time
+	 * in the likely event that we reset or delete the context later.
+	 */
+	while ((cb = context->reset_cbs) != NULL)
+	{
+		context->reset_cbs = cb->next;
+		(*cb->func) (cb->arg);
+	}
+}
+
 /*
  * MemoryContextSetParent
  *		Change a context to belong to a new parent (or no parent).
@@ -318,9 +377,8 @@ void
 MemoryContextAllowInCriticalSection(MemoryContext context, bool allow)
 {
 	AssertArg(MemoryContextIsValid(context));
-#ifdef USE_ASSERT_CHECKING
+
 	context->allowInCritSection = allow;
-#endif
 }
 
 /*
@@ -589,11 +647,8 @@ MemoryContextCreate(NodeTag tag, Size size,
 		node->parent = parent;
 		node->nextchild = parent->firstchild;
 		parent->firstchild = node;
-
-#ifdef USE_ASSERT_CHECKING
 		/* inherit allowInCritSection flag from parent */
 		node->allowInCritSection = parent->allowInCritSection;
-#endif
 	}
 
 	VALGRIND_CREATE_MEMPOOL(node, 0, false);

src/include/nodes/memnodes.h

@@ -16,6 +16,22 @@
 
 #include "nodes/nodes.h"
 
+/*
+ * A memory context can have callback functions registered on it.  Any such
+ * function will be called once just before the context is next reset or
+ * deleted.  The MemoryContextCallback struct describing such a callback
+ * typically would be allocated within the context itself, thereby avoiding
+ * any need to manage it explicitly (the reset/delete action will free it).
+ */
+typedef void (*MemoryContextCallbackFunction) (void *arg);
+
+typedef struct MemoryContextCallback
+{
+	MemoryContextCallbackFunction func; /* function to call */
+	void	   *arg;			/* argument to pass it */
+	struct MemoryContextCallback *next; /* next in list of callbacks */
+} MemoryContextCallback;
+
 /*
  * MemoryContext
  *		A logical context in which memory allocations occur.
@@ -54,15 +70,15 @@ typedef struct MemoryContextMethods
 typedef struct MemoryContextData
 {
 	NodeTag		type;			/* identifies exact kind of context */
+	/* these two fields are placed here to minimize alignment wastage: */
+	bool		isReset;		/* T = no space alloced since last reset */
+	bool		allowInCritSection;		/* allow palloc in critical section */
 	MemoryContextMethods *methods;		/* virtual function table */
 	MemoryContext parent;		/* NULL if no parent (toplevel context) */
 	MemoryContext firstchild;	/* head of linked list of children */
 	MemoryContext nextchild;	/* next child of same parent */
 	char	   *name;			/* context name (just for debugging) */
-	bool		isReset;		/* T = no space alloced since last reset */
-#ifdef USE_ASSERT_CHECKING
-	bool		allowInCritSection;	/* allow palloc in critical section */
-#endif
+	MemoryContextCallback *reset_cbs;	/* list of reset/delete callbacks */
 } MemoryContextData;
 
 /* utils/palloc.h contains typedef struct MemoryContextData *MemoryContext */

src/include/utils/memutils.h

@@ -94,6 +94,8 @@ extern void MemoryContextDelete(MemoryContext context);
 extern void MemoryContextResetChildren(MemoryContext context);
 extern void MemoryContextDeleteChildren(MemoryContext context);
 extern void MemoryContextResetAndDeleteChildren(MemoryContext context);
+extern void MemoryContextRegisterResetCallback(MemoryContext context,
+								   MemoryContextCallback *cb);
 extern void MemoryContextSetParent(MemoryContext context,
 					   MemoryContext new_parent);
 extern Size GetMemoryChunkSpace(void *pointer);