Re: Transaction-lifespan memory leak with plpgsql DO blocks

Robert Haas <robertmhaas@gmail.com> writes:
> On Tue, Nov 12, 2013 at 11:18 AM, Tom Lane <tgl@sss.pgh.pa.us> wrote:
>> Or we could say "what the heck are you doing executing tens of
>> thousands of DO blocks?  Make it into a real live function;
>> you'll save a lot of cycles on parsing costs."  I'm not sure that
>> this is a usage pattern we ought to be optimizing for.

> I'm not volunteering to spend time fixing this, but I disagree with
> the premise that it isn't worth fixing, because I think it's a POLA
> violation.

Yeah, I'm not terribly comfortable with letting it go either.  Attached
is a proposed patch.  I couldn't see any nice way to do it without adding
a field to PLpgSQL_execstate, so this isn't a feasible solution for
back-patching (it'd break the plpgsql debugger).  However, given the
infrequency of complaints, I think fixing it in 9.4 and up is good enough.

I checked that this eliminates the memory leak using this test case:

do $outer$
begin
  for i in 1..1000000 loop
    execute $e$
      do $$
      declare x int = 0;
      begin
        x := x + 1;
      end;
      $$;
    $e$;
  end loop;
end;
$outer$;

which eats a couple GB in HEAD and nothing with the patch.
The run time seems to be the same or a bit less, too.

Any objections to applying this to HEAD?

            regards, tom lane

diff --git a/src/pl/plpgsql/src/pl_exec.c b/src/pl/plpgsql/src/pl_exec.c
index bc31fe9..76da842 100644
*** a/src/pl/plpgsql/src/pl_exec.c
--- b/src/pl/plpgsql/src/pl_exec.c
*************** typedef struct
*** 50,81 ****
      bool       *freevals;        /* which arguments are pfree-able */
  } PreparedParamsData;

! /*
!  * All plpgsql function executions within a single transaction share the same
!  * executor EState for evaluating "simple" expressions.  Each function call
!  * creates its own "eval_econtext" ExprContext within this estate for
!  * per-evaluation workspace.  eval_econtext is freed at normal function exit,
!  * and the EState is freed at transaction end (in case of error, we assume
!  * that the abort mechanisms clean it all up).    Furthermore, any exception
!  * block within a function has to have its own eval_econtext separate from
!  * the containing function's, so that we can clean up ExprContext callbacks
!  * properly at subtransaction exit.  We maintain a stack that tracks the
!  * individual econtexts so that we can clean up correctly at subxact exit.
!  *
!  * This arrangement is a bit tedious to maintain, but it's worth the trouble
!  * so that we don't have to re-prepare simple expressions on each trip through
!  * a function.    (We assume the case to optimize is many repetitions of a
!  * function within a transaction.)
!  */
! typedef struct SimpleEcontextStackEntry
! {
!     ExprContext *stack_econtext;    /* a stacked econtext */
!     SubTransactionId xact_subxid;        /* ID for current subxact */
!     struct SimpleEcontextStackEntry *next;        /* next stack entry up */
! } SimpleEcontextStackEntry;
!
! static EState *simple_eval_estate = NULL;
! static SimpleEcontextStackEntry *simple_econtext_stack = NULL;

  /************************************************************
   * Local function forward declarations
--- 50,57 ----
      bool       *freevals;        /* which arguments are pfree-able */
  } PreparedParamsData;

! /* Shared simple-expression eval context for regular plpgsql functions */
! static SimpleEvalContext simple_eval_context = {NULL, NULL};

  /************************************************************
   * Local function forward declarations
*************** static int exec_stmt_dynfors(PLpgSQL_exe
*** 136,142 ****

  static void plpgsql_estate_setup(PLpgSQL_execstate *estate,
                       PLpgSQL_function *func,
!                      ReturnSetInfo *rsi);
  static void exec_eval_cleanup(PLpgSQL_execstate *estate);

  static void exec_prepare_plan(PLpgSQL_execstate *estate,
--- 112,119 ----

  static void plpgsql_estate_setup(PLpgSQL_execstate *estate,
                       PLpgSQL_function *func,
!                      ReturnSetInfo *rsi,
!                      SimpleEvalContext *simple_context);
  static void exec_eval_cleanup(PLpgSQL_execstate *estate);

  static void exec_prepare_plan(PLpgSQL_execstate *estate,
*************** static char *format_preparedparamsdata(P
*** 230,239 ****
  /* ----------
   * plpgsql_exec_function    Called by the call handler for
   *                function execution.
   * ----------
   */
  Datum
! plpgsql_exec_function(PLpgSQL_function *func, FunctionCallInfo fcinfo)
  {
      PLpgSQL_execstate estate;
      ErrorContextCallback plerrcontext;
--- 207,222 ----
  /* ----------
   * plpgsql_exec_function    Called by the call handler for
   *                function execution.
+  *
+  * This is also used to execute inline code blocks (DO blocks).  The only
+  * difference that this code is aware of is that for a DO block, we want
+  * to use a private SimpleEvalContext, whose address must be passed as
+  * simple_context.  For regular functions, pass NULL.
   * ----------
   */
  Datum
! plpgsql_exec_function(PLpgSQL_function *func, FunctionCallInfo fcinfo,
!                       SimpleEvalContext *simple_context)
  {
      PLpgSQL_execstate estate;
      ErrorContextCallback plerrcontext;
*************** plpgsql_exec_function(PLpgSQL_function *
*** 243,249 ****
      /*
       * Setup the execution state
       */
!     plpgsql_estate_setup(&estate, func, (ReturnSetInfo *) fcinfo->resultinfo);

      /*
       * Setup error traceback support for ereport()
--- 226,233 ----
      /*
       * Setup the execution state
       */
!     plpgsql_estate_setup(&estate, func, (ReturnSetInfo *) fcinfo->resultinfo,
!                          simple_context);

      /*
       * Setup error traceback support for ereport()
*************** plpgsql_exec_trigger(PLpgSQL_function *f
*** 503,509 ****
      /*
       * Setup the execution state
       */
!     plpgsql_estate_setup(&estate, func, NULL);

      /*
       * Setup error traceback support for ereport()
--- 487,493 ----
      /*
       * Setup the execution state
       */
!     plpgsql_estate_setup(&estate, func, NULL, NULL);

      /*
       * Setup error traceback support for ereport()
*************** plpgsql_exec_event_trigger(PLpgSQL_funct
*** 782,788 ****
      /*
       * Setup the execution state
       */
!     plpgsql_estate_setup(&estate, func, NULL);

      /*
       * Setup error traceback support for ereport()
--- 766,772 ----
      /*
       * Setup the execution state
       */
!     plpgsql_estate_setup(&estate, func, NULL, NULL);

      /*
       * Setup error traceback support for ereport()
*************** exec_stmt_raise(PLpgSQL_execstate *estat
*** 3087,3093 ****
  static void
  plpgsql_estate_setup(PLpgSQL_execstate *estate,
                       PLpgSQL_function *func,
!                      ReturnSetInfo *rsi)
  {
      /* this link will be restored at exit from plpgsql_call_handler */
      func->cur_estate = estate;
--- 3071,3078 ----
  static void
  plpgsql_estate_setup(PLpgSQL_execstate *estate,
                       PLpgSQL_function *func,
!                      ReturnSetInfo *rsi,
!                      SimpleEvalContext *simple_context)
  {
      /* this link will be restored at exit from plpgsql_call_handler */
      func->cur_estate = estate;
*************** plpgsql_estate_setup(PLpgSQL_execstate *
*** 3126,3131 ****
--- 3111,3122 ----
      estate->datums = palloc(sizeof(PLpgSQL_datum *) * estate->ndatums);
      /* caller is expected to fill the datums array */

+     /* set up for use of appropriate simple-expression context */
+     if (simple_context)
+         estate->simple_context = simple_context;
+     else
+         estate->simple_context = &simple_eval_context;
+
      estate->eval_tuptable = NULL;
      estate->eval_processed = 0;
      estate->eval_lastoid = InvalidOid;
*************** exec_eval_simple_expr(PLpgSQL_execstate
*** 5148,5154 ****
       */
      if (expr->expr_simple_lxid != curlxid)
      {
!         oldcontext = MemoryContextSwitchTo(simple_eval_estate->es_query_cxt);
          expr->expr_simple_state = ExecInitExpr(expr->expr_simple_expr, NULL);
          expr->expr_simple_in_use = false;
          expr->expr_simple_lxid = curlxid;
--- 5139,5145 ----
       */
      if (expr->expr_simple_lxid != curlxid)
      {
!         oldcontext = MemoryContextSwitchTo(estate->simple_context->estate->es_query_cxt);
          expr->expr_simple_state = ExecInitExpr(expr->expr_simple_expr, NULL);
          expr->expr_simple_in_use = false;
          expr->expr_simple_lxid = curlxid;
*************** exec_set_found(PLpgSQL_execstate *estate
*** 6190,6196 ****
  /*
   * plpgsql_create_econtext --- create an eval_econtext for the current function
   *
!  * We may need to create a new simple_eval_estate too, if there's not one
   * already for the current transaction.  The EState will be cleaned up at
   * transaction end.
   */
--- 6181,6187 ----
  /*
   * plpgsql_create_econtext --- create an eval_econtext for the current function
   *
!  * We may need to create a new simple-eval EState too, if there's not one
   * already for the current transaction.  The EState will be cleaned up at
   * transaction end.
   */
*************** plpgsql_create_econtext(PLpgSQL_execstat
*** 6204,6222 ****
       * one already in the current transaction.    The EState is made a child of
       * TopTransactionContext so it will have the right lifespan.
       */
!     if (simple_eval_estate == NULL)
      {
          MemoryContext oldcontext;

          oldcontext = MemoryContextSwitchTo(TopTransactionContext);
!         simple_eval_estate = CreateExecutorState();
          MemoryContextSwitchTo(oldcontext);
      }

      /*
       * Create a child econtext for the current function.
       */
!     estate->eval_econtext = CreateExprContext(simple_eval_estate);

      /*
       * Make a stack entry so we can clean up the econtext at subxact end.
--- 6195,6213 ----
       * one already in the current transaction.    The EState is made a child of
       * TopTransactionContext so it will have the right lifespan.
       */
!     if (estate->simple_context->estate == NULL)
      {
          MemoryContext oldcontext;

          oldcontext = MemoryContextSwitchTo(TopTransactionContext);
!         estate->simple_context->estate = CreateExecutorState();
          MemoryContextSwitchTo(oldcontext);
      }

      /*
       * Create a child econtext for the current function.
       */
!     estate->eval_econtext = CreateExprContext(estate->simple_context->estate);

      /*
       * Make a stack entry so we can clean up the econtext at subxact end.
*************** plpgsql_create_econtext(PLpgSQL_execstat
*** 6229,6236 ****
      entry->stack_econtext = estate->eval_econtext;
      entry->xact_subxid = GetCurrentSubTransactionId();

!     entry->next = simple_econtext_stack;
!     simple_econtext_stack = entry;
  }

  /*
--- 6220,6227 ----
      entry->stack_econtext = estate->eval_econtext;
      entry->xact_subxid = GetCurrentSubTransactionId();

!     entry->next = estate->simple_context->stack;
!     estate->simple_context->stack = entry;
  }

  /*
*************** plpgsql_destroy_econtext(PLpgSQL_execsta
*** 6244,6255 ****
  {
      SimpleEcontextStackEntry *next;

!     Assert(simple_econtext_stack != NULL);
!     Assert(simple_econtext_stack->stack_econtext == estate->eval_econtext);

!     next = simple_econtext_stack->next;
!     pfree(simple_econtext_stack);
!     simple_econtext_stack = next;

      FreeExprContext(estate->eval_econtext, true);
      estate->eval_econtext = NULL;
--- 6235,6246 ----
  {
      SimpleEcontextStackEntry *next;

!     Assert(estate->simple_context->stack != NULL);
!     Assert(estate->simple_context->stack->stack_econtext == estate->eval_econtext);

!     next = estate->simple_context->stack->next;
!     pfree(estate->simple_context->stack);
!     estate->simple_context->stack = next;

      FreeExprContext(estate->eval_econtext, true);
      estate->eval_econtext = NULL;
*************** plpgsql_destroy_econtext(PLpgSQL_execsta
*** 6258,6264 ****
  /*
   * plpgsql_xact_cb --- post-transaction-commit-or-abort cleanup
   *
!  * If a simple-expression EState was created in the current transaction,
   * it has to be cleaned up.
   */
  void
--- 6249,6255 ----
  /*
   * plpgsql_xact_cb --- post-transaction-commit-or-abort cleanup
   *
!  * If a shared simple-expression EState was created in the current transaction,
   * it has to be cleaned up.
   */
  void
*************** plpgsql_xact_cb(XactEvent event, void *a
*** 6273,6288 ****
      if (event == XACT_EVENT_COMMIT || event == XACT_EVENT_PREPARE)
      {
          /* Shouldn't be any econtext stack entries left at commit */
!         Assert(simple_econtext_stack == NULL);

!         if (simple_eval_estate)
!             FreeExecutorState(simple_eval_estate);
!         simple_eval_estate = NULL;
      }
      else if (event == XACT_EVENT_ABORT)
      {
!         simple_econtext_stack = NULL;
!         simple_eval_estate = NULL;
      }
  }

--- 6264,6279 ----
      if (event == XACT_EVENT_COMMIT || event == XACT_EVENT_PREPARE)
      {
          /* Shouldn't be any econtext stack entries left at commit */
!         Assert(simple_eval_context.stack == NULL);

!         if (simple_eval_context.estate)
!             FreeExecutorState(simple_eval_context.estate);
!         simple_eval_context.estate = NULL;
      }
      else if (event == XACT_EVENT_ABORT)
      {
!         simple_eval_context.stack = NULL;
!         simple_eval_context.estate = NULL;
      }
  }

*************** plpgsql_xact_cb(XactEvent event, void *a
*** 6291,6297 ****
   *
   * Make sure any simple-expression econtexts created in the current
   * subtransaction get cleaned up.  We have to do this explicitly because
!  * no other code knows which child econtexts of simple_eval_estate belong
   * to which level of subxact.
   */
  void
--- 6282,6288 ----
   *
   * Make sure any simple-expression econtexts created in the current
   * subtransaction get cleaned up.  We have to do this explicitly because
!  * no other code knows which child econtexts of the simple-eval EState belong
   * to which level of subxact.
   */
  void
*************** plpgsql_subxact_cb(SubXactEvent event, S
*** 6300,6315 ****
  {
      if (event == SUBXACT_EVENT_COMMIT_SUB || event == SUBXACT_EVENT_ABORT_SUB)
      {
!         while (simple_econtext_stack != NULL &&
!                simple_econtext_stack->xact_subxid == mySubid)
          {
              SimpleEcontextStackEntry *next;

!             FreeExprContext(simple_econtext_stack->stack_econtext,
                              (event == SUBXACT_EVENT_COMMIT_SUB));
!             next = simple_econtext_stack->next;
!             pfree(simple_econtext_stack);
!             simple_econtext_stack = next;
          }
      }
  }
--- 6291,6306 ----
  {
      if (event == SUBXACT_EVENT_COMMIT_SUB || event == SUBXACT_EVENT_ABORT_SUB)
      {
!         while (simple_eval_context.stack != NULL &&
!                simple_eval_context.stack->xact_subxid == mySubid)
          {
              SimpleEcontextStackEntry *next;

!             FreeExprContext(simple_eval_context.stack->stack_econtext,
                              (event == SUBXACT_EVENT_COMMIT_SUB));
!             next = simple_eval_context.stack->next;
!             pfree(simple_eval_context.stack);
!             simple_eval_context.stack = next;
          }
      }
  }
diff --git a/src/pl/plpgsql/src/pl_handler.c b/src/pl/plpgsql/src/pl_handler.c
index 912422c..add8321 100644
*** a/src/pl/plpgsql/src/pl_handler.c
--- b/src/pl/plpgsql/src/pl_handler.c
*************** plpgsql_call_handler(PG_FUNCTION_ARGS)
*** 136,142 ****
              retval = (Datum) 0;
          }
          else
!             retval = plpgsql_exec_function(func, fcinfo);
      }
      PG_CATCH();
      {
--- 136,142 ----
              retval = (Datum) 0;
          }
          else
!             retval = plpgsql_exec_function(func, fcinfo, NULL);
      }
      PG_CATCH();
      {
*************** plpgsql_inline_handler(PG_FUNCTION_ARGS)
*** 175,180 ****
--- 175,181 ----
      PLpgSQL_function *func;
      FunctionCallInfoData fake_fcinfo;
      FmgrInfo    flinfo;
+     SimpleEvalContext simple_context;
      Datum        retval;
      int            rc;

*************** plpgsql_inline_handler(PG_FUNCTION_ARGS)
*** 203,209 ****
      flinfo.fn_oid = InvalidOid;
      flinfo.fn_mcxt = CurrentMemoryContext;

!     retval = plpgsql_exec_function(func, &fake_fcinfo);

      /* Function should now have no remaining use-counts ... */
      func->use_count--;
--- 204,235 ----
      flinfo.fn_oid = InvalidOid;
      flinfo.fn_mcxt = CurrentMemoryContext;

!     /* Initialize private context for simple-expression evaluation */
!     simple_context.estate = NULL;
!     simple_context.stack = NULL;
!
!     /* Call the handler */
!     PG_TRY();
!     {
!         retval = plpgsql_exec_function(func, &fake_fcinfo, &simple_context);
!     }
!     PG_CATCH();
!     {
!         /* Release private estate, if we got as far as making one */
!         if (simple_context.estate)
!             FreeExecutorState(simple_context.estate);
!
!         /* And propagate the error */
!         PG_RE_THROW();
!     }
!     PG_END_TRY();
!
!     /* Shouldn't be any econtext stack entries left, if no error */
!     Assert(simple_context.stack == NULL);
!
!     /* Clean up the private estate, if any */
!     if (simple_context.estate)
!         FreeExecutorState(simple_context.estate);

      /* Function should now have no remaining use-counts ... */
      func->use_count--;
diff --git a/src/pl/plpgsql/src/plpgsql.h b/src/pl/plpgsql/src/plpgsql.h
index 9cb4f53..29c13d1 100644
*** a/src/pl/plpgsql/src/plpgsql.h
--- b/src/pl/plpgsql/src/plpgsql.h
*************** typedef enum
*** 166,171 ****
--- 166,212 ----
  } PLpgSQL_resolve_option;


+ /*
+  * All plpgsql function executions within a single transaction share the same
+  * executor EState for evaluating "simple" expressions.  Each function call
+  * creates its own "eval_econtext" ExprContext within this estate for
+  * per-evaluation workspace.  eval_econtext is freed at normal function exit,
+  * and the EState is freed at transaction end (in case of error, we assume
+  * that the abort mechanisms clean it all up).    Furthermore, any exception
+  * block within a function has to have its own eval_econtext separate from
+  * the containing function's, so that we can clean up ExprContext callbacks
+  * properly at subtransaction exit.  We maintain a stack that tracks the
+  * individual econtexts so that we can clean up correctly at subxact exit.
+  *
+  * This arrangement is a bit tedious to maintain, but it's worth the trouble
+  * so that we don't have to re-prepare simple expressions on each trip through
+  * a function.    (We assume the case to optimize is many repetitions of a
+  * function within a transaction.)
+  *
+  * However, there's no value in trying to amortize simple expression setup
+  * across multiple executions of a DO block (inline code block), since there
+  * can never be any.  If we use the shared EState for a DO block, the expr
+  * state trees are effectively leaked till end of transaction, and that can
+  * add up if the user keeps on submitting DO blocks.  Therefore, each DO block
+  * has its own simple-expression EState, which is cleaned up at exit from
+  * plpgsql_inline_handler().  To minimize the notational difference between
+  * function and DO block execution, code should access the appropriate state
+  * via PLpgSQL_execstate's simple_context pointer.
+  */
+ typedef struct SimpleEcontextStackEntry
+ {
+     ExprContext *stack_econtext;    /* a stacked econtext */
+     SubTransactionId xact_subxid;        /* ID for current subxact */
+     struct SimpleEcontextStackEntry *next;        /* next stack entry up */
+ } SimpleEcontextStackEntry;
+
+ typedef struct SimpleEvalContext
+ {
+     EState       *estate;            /* EState to use (NULL if not made yet) */
+     SimpleEcontextStackEntry *stack;    /* top of stack of ExprContexts */
+ } SimpleEvalContext;
+
+
  /**********************************************************************
   * Node and structure definitions
   **********************************************************************/
*************** typedef struct PLpgSQL_execstate
*** 773,782 ****
--- 814,827 ----
      ResourceOwner tuple_store_owner;
      ReturnSetInfo *rsi;

+     /* the datums representing the function's local variables */
      int            found_varno;
      int            ndatums;
      PLpgSQL_datum **datums;

+     /* state to use for "simple" expression evaluation */
+     SimpleEvalContext *simple_context;
+
      /* temporary state for results from evaluation of query or expr */
      SPITupleTable *eval_tuptable;
      uint32        eval_processed;
*************** extern Datum plpgsql_validator(PG_FUNCTI
*** 943,949 ****
   * ----------
   */
  extern Datum plpgsql_exec_function(PLpgSQL_function *func,
!                       FunctionCallInfo fcinfo);
  extern HeapTuple plpgsql_exec_trigger(PLpgSQL_function *func,
                       TriggerData *trigdata);
  extern void plpgsql_exec_event_trigger(PLpgSQL_function *func,
--- 988,995 ----
   * ----------
   */
  extern Datum plpgsql_exec_function(PLpgSQL_function *func,
!                       FunctionCallInfo fcinfo,
!                       SimpleEvalContext *simple_context);
  extern HeapTuple plpgsql_exec_trigger(PLpgSQL_function *func,
                       TriggerData *trigdata);
  extern void plpgsql_exec_event_trigger(PLpgSQL_function *func,