diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index 202bc785bcfd2001088e201a36e369968cd1f7d7..72a32abdddb671769d8999a3f30b009ea603b900 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/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);
diff --git a/src/include/nodes/memnodes.h b/src/include/nodes/memnodes.h
index ca9c3dea3cb91e1557057309cb4acc3189c9175f..3eeaad49288e88a02892e691dc0cf24961dcc5f8 100644
--- a/src/include/nodes/memnodes.h
+++ b/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 */
diff --git a/src/include/utils/memutils.h b/src/include/utils/memutils.h
index 85aba7a7220b9eeda39126b52798f61c824b998a..76ad7d443b03053f4637467e72d2f1bb1fc4f26a 100644
--- a/src/include/utils/memutils.h
+++ b/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);