MemoryContext reset/delete callbacks - Mailing list pgsql-hackers

We discussed this idea a couple weeks ago.  The core of it is that when a
memory context is being deleted, you might want something extra to happen
beyond just pfree'ing everything in the context.  I'm thinking in
particular that this might provide a nice solution to the problem we
discussed earlier today of managing cached information about a domain's
CHECK constraints: a function could make a reference-counted link to some
data in its FmgrInfo cache, and use a memory context callback on the
context containing the FmgrInfo to ensure that the reference count gets
decremented when needed and not before.

Attached is a draft patch for this.  I've not really tested it more than
seeing that it compiles, but it's so simple that there are unlikely to be
bugs as such in it.  There are some design decisions that could be
questioned though:

1. I used ilists for the linked list of callback requests.  This creates a
dependency from memory contexts to lib/ilist.  That's all right at the
moment since lib/ilist does no memory allocation, but it might be
logically problematic.  We could just use explicit "struct foo *" links
with little if any notational penalty, so I wonder if that would be
better.

2. I specified that callers have to allocate the storage for the callback
structs.  This saves a palloc cycle in just about every use-case I've
thought of, since callers generally will need to palloc some larger struct
of their own and they can just embed the MemoryContextCallback struct in
that.  It does seem like this offers a larger surface for screwups, in
particular putting the callback struct in the wrong memory context --- but
there's plenty of surface for that type of mistake anyway, if you put
whatever the "void *arg" is pointing at in the wrong context.
So I'm OK with this but could also accept a design in which
MemoryContextRegisterResetCallback takes just a function pointer and a
"void *arg" and palloc's the callback struct for itself in the target
context.  I'm unsure whether that adds enough safety to justify a separate
palloc.

3. Is the "void *arg" API for the callback func sufficient?  I considered
also passing it the MemoryContext, but couldn't really find a use-case
to justify that.

4. I intentionally didn't provide a MemoryContextDeregisterResetCallback
API.  Since it's just a singly-linked list, that could be an expensive
operation and so I'd rather discourage it.  In the use cases I've thought
of, you'd want the callback to remain active for the life of the context
anyway, so there would be no need.  (And, of course, there is slist_delete
for the truly determined ...)

Comments?

            regards, tom lane

diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index 202bc78..624d08a 100644
*** a/src/backend/utils/mmgr/mcxt.c
--- b/src/backend/utils/mmgr/mcxt.c
*************** MemoryContext CurTransactionContext = NU
*** 54,59 ****
--- 54,60 ----
  /* 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);

  /*
*************** MemoryContextReset(MemoryContext context
*** 150,155 ****
--- 151,157 ----
      /* 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);
*************** MemoryContextDelete(MemoryContext contex
*** 196,201 ****
--- 198,211 ----
      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
       * to the context tree.  Better a leak than a crash.
*************** MemoryContextResetAndDeleteChildren(Memo
*** 243,248 ****
--- 253,308 ----
  }

  /*
+  * 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. */
+     slist_push_head(&context->reset_callbacks, &cb->list_node);
+     /* 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)
+ {
+     /*
+      * 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 (!slist_is_empty(&context->reset_callbacks))
+     {
+         MemoryContextCallback *cb;
+
+         cb = slist_container(MemoryContextCallback, list_node,
+                              slist_pop_head_node(&context->reset_callbacks));
+         (*cb->func) (cb->arg);
+     }
+ }
+
+ /*
   * MemoryContextSetParent
   *        Change a context to belong to a new parent (or no parent).
   *
diff --git a/src/include/nodes/memnodes.h b/src/include/nodes/memnodes.h
index ca9c3de..2440f02 100644
*** a/src/include/nodes/memnodes.h
--- b/src/include/nodes/memnodes.h
***************
*** 14,19 ****
--- 14,20 ----
  #ifndef MEMNODES_H
  #define MEMNODES_H

+ #include "lib/ilist.h"
  #include "nodes/nodes.h"

  /*
*************** typedef struct MemoryContextData
*** 59,72 ****
      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
  } MemoryContextData;

  /* utils/palloc.h contains typedef struct MemoryContextData *MemoryContext */


  /*
   * MemoryContextIsValid
--- 60,90 ----
      MemoryContext firstchild;    /* head of linked list of children */
      MemoryContext nextchild;    /* next child of same parent */
      char       *name;            /* context name (just for debugging) */
+     slist_head    reset_callbacks;    /* callbacks for context reset/delete */
      bool        isReset;        /* T = no space alloced since last reset */
  #ifdef USE_ASSERT_CHECKING
!     bool        allowInCritSection;        /* allow palloc in critical section */
  #endif
  } MemoryContextData;

  /* utils/palloc.h contains typedef struct MemoryContextData *MemoryContext */

+ /*
+  * 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 */
+     slist_node    list_node;        /* link in linked list of callbacks */
+ } MemoryContextCallback;
+

  /*
   * MemoryContextIsValid
diff --git a/src/include/utils/memutils.h b/src/include/utils/memutils.h
index 85aba7a..76ad7d4 100644
*** a/src/include/utils/memutils.h
--- b/src/include/utils/memutils.h
*************** extern void MemoryContextDelete(MemoryCo
*** 94,99 ****
--- 94,101 ----
  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);