Re: Transaction-lifespan memory leak with plpgsql DO blocks

I wrote:
> Robert Haas <robertmhaas@gmail.com> writes:
>> 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.

This patch crashed and burned when I tried it on a case where a DO block
traps an exception :-(.  I had thought that a private econtext stack was
the right thing to do, but it isn't because we need plpgsql_subxact_cb
to be able to clean up dead econtexts in either functions or DO blocks.
So after some experimentation I came up with version 2, attached.
The additional test case I was using looks like

do $outer$
begin
  for i in 1..100000 loop
   begin
    execute $e$
      do $$
      declare x int = 0;
      begin
        x := 1 / x;
      end;
      $$;
    $e$;
  exception when division_by_zero then null;
  end;
  end loop;
end;
$outer$;

If you try this with the patch, you'll notice that there's still a slow
leak, but that's not the fault of DO blocks: the same leak exists if you
transpose this code into a regular function.  That leak is the fault of
exec_stmt_dynexecute, which copies the querystring into the function's
main execution context, from which it won't get freed if an error is
thrown out of SPI_execute.  (If we were using EXECUTE USING, the "ppd"
structure would get leaked too.)  There are some other similar error-
case leaks in other plpgsql statements, I think.  I'm not excited
about trying to clean those up as part of this patch.

            regards, tom lane

diff --git a/src/pl/plpgsql/src/pl_exec.c b/src/pl/plpgsql/src/pl_exec.c
index bc31fe9..f5f1892 100644
*** a/src/pl/plpgsql/src/pl_exec.c
--- b/src/pl/plpgsql/src/pl_exec.c
*************** typedef struct
*** 66,71 ****
--- 66,80 ----
   * 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().  DO blocks still use the simple_econtext_stack,
+  * though, so that subxact abort cleanup does the right thing.
   */
  typedef struct SimpleEcontextStackEntry
  {
*************** typedef struct SimpleEcontextStackEntry
*** 74,80 ****
      struct SimpleEcontextStackEntry *next;        /* next stack entry up */
  } SimpleEcontextStackEntry;

! static EState *simple_eval_estate = NULL;
  static SimpleEcontextStackEntry *simple_econtext_stack = NULL;

  /************************************************************
--- 83,89 ----
      struct SimpleEcontextStackEntry *next;        /* next stack entry up */
  } SimpleEcontextStackEntry;

! static EState *shared_simple_eval_estate = NULL;
  static SimpleEcontextStackEntry *simple_econtext_stack = NULL;

  /************************************************************
*************** 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,
--- 145,152 ----

  static void plpgsql_estate_setup(PLpgSQL_execstate *estate,
                       PLpgSQL_function *func,
!                      ReturnSetInfo *rsi,
!                      EState *simple_eval_estate);
  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;
--- 240,256 ----
  /* ----------
   * 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 simple_eval_estate, which is created and passed in by
+  * the caller.  For regular functions, pass NULL, which implies using
+  * shared_simple_eval_estate.
   * ----------
   */
  Datum
! plpgsql_exec_function(PLpgSQL_function *func, FunctionCallInfo fcinfo,
!                       EState *simple_eval_estate)
  {
      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()
--- 260,267 ----
      /*
       * Setup the execution state
       */
!     plpgsql_estate_setup(&estate, func, (ReturnSetInfo *) fcinfo->resultinfo,
!                          simple_eval_estate);

      /*
       * 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()
--- 521,527 ----
      /*
       * 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()
--- 800,806 ----
      /*
       * 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;
--- 3105,3112 ----
  static void
  plpgsql_estate_setup(PLpgSQL_execstate *estate,
                       PLpgSQL_function *func,
!                      ReturnSetInfo *rsi,
!                      EState *simple_eval_estate)
  {
      /* this link will be restored at exit from plpgsql_call_handler */
      func->cur_estate = estate;
*************** plpgsql_estate_setup(PLpgSQL_execstate *
*** 3126,3131 ****
--- 3145,3156 ----
      estate->datums = palloc(sizeof(PLpgSQL_datum *) * estate->ndatums);
      /* caller is expected to fill the datums array */

+     /* set up for use of appropriate simple-expression EState */
+     if (simple_eval_estate)
+         estate->simple_eval_estate = simple_eval_estate;
+     else
+         estate->simple_eval_estate = shared_simple_eval_estate;
+
      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;
--- 5173,5179 ----
       */
      if (expr->expr_simple_lxid != curlxid)
      {
!         oldcontext = MemoryContextSwitchTo(estate->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;
*************** exec_set_found(PLpgSQL_execstate *estate
*** 6190,6197 ****
  /*
   * 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.
   */
  static void
--- 6215,6222 ----
  /*
   * plpgsql_create_econtext --- create an eval_econtext for the current function
   *
!  * We may need to create a new shared_simple_eval_estate too, if there's not
!  * one already for the current transaction.  The EState will be cleaned up at
   * transaction end.
   */
  static void
*************** plpgsql_create_econtext(PLpgSQL_execstat
*** 6203,6222 ****
       * Create an EState for evaluation of simple expressions, if there's not
       * 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.
--- 6228,6252 ----
       * Create an EState for evaluation of simple expressions, if there's not
       * one already in the current transaction.    The EState is made a child of
       * TopTransactionContext so it will have the right lifespan.
+      *
+      * Note that this path is never taken when executing a DO block; the
+      * required EState was already made by plpgsql_inline_handler.
       */
!     if (estate->simple_eval_estate == NULL)
      {
          MemoryContext oldcontext;

+         Assert(shared_simple_eval_estate == NULL);
          oldcontext = MemoryContextSwitchTo(TopTransactionContext);
!         shared_simple_eval_estate = CreateExecutorState();
!         estate->simple_eval_estate = shared_simple_eval_estate;
          MemoryContextSwitchTo(oldcontext);
      }

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

      /*
       * Make a stack entry so we can clean up the econtext at subxact end.
*************** plpgsql_xact_cb(XactEvent event, void *a
*** 6275,6288 ****
          /* 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;
      }
  }

--- 6305,6318 ----
          /* Shouldn't be any econtext stack entries left at commit */
          Assert(simple_econtext_stack == NULL);

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

*************** plpgsql_xact_cb(XactEvent event, void *a
*** 6291,6298 ****
   *
   * 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
  plpgsql_subxact_cb(SubXactEvent event, SubTransactionId mySubid,
--- 6321,6327 ----
   *
   * 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 econtexts belong to which level of subxact.
   */
  void
  plpgsql_subxact_cb(SubXactEvent event, SubTransactionId mySubid,
diff --git a/src/pl/plpgsql/src/pl_handler.c b/src/pl/plpgsql/src/pl_handler.c
index 912422c..5bfe3c3 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;
+     EState       *simple_eval_estate;
      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,254 ----
      flinfo.fn_oid = InvalidOid;
      flinfo.fn_mcxt = CurrentMemoryContext;

!     /* Create a private EState for simple-expression execution */
!     simple_eval_estate = CreateExecutorState();
!
!     /* And run the function */
!     PG_TRY();
!     {
!         retval = plpgsql_exec_function(func, &fake_fcinfo, simple_eval_estate);
!     }
!     PG_CATCH();
!     {
!         /*
!          * We need to clean up what would otherwise be long-lived resources
!          * accumulated by the failed DO block, principally cached plans for
!          * statements (which can be flushed with plpgsql_free_function_memory)
!          * and execution trees for simple expressions, which are in the
!          * private EState.
!          *
!          * Before releasing the private EState, we must clean up any
!          * simple_econtext_stack entries pointing into it, which we can do by
!          * invoking the subxact callback.  (It will be called again later if
!          * some outer control level does a subtransaction abort, but no harm
!          * is done.)  We cheat a bit knowing that plpgsql_subxact_cb does not
!          * pay attention to its parentSubid argument.
!          */
!         plpgsql_subxact_cb(SUBXACT_EVENT_ABORT_SUB,
!                            GetCurrentSubTransactionId(),
!                            0, NULL);
!
!         /* Clean up the private EState */
!         FreeExecutorState(simple_eval_estate);
!
!         /* Function should now have no remaining use-counts ... */
!         func->use_count--;
!         Assert(func->use_count == 0);
!
!         /* ... so we can free subsidiary storage */
!         plpgsql_free_function_memory(func);
!
!         /* And propagate the error */
!         PG_RE_THROW();
!     }
!     PG_END_TRY();
!
!     /* Clean up the private EState */
!     FreeExecutorState(simple_eval_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..3afcdf3 100644
*** a/src/pl/plpgsql/src/plpgsql.h
--- b/src/pl/plpgsql/src/plpgsql.h
*************** typedef struct PLpgSQL_execstate
*** 773,782 ****
--- 773,786 ----
      ResourceOwner tuple_store_owner;
      ReturnSetInfo *rsi;

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

+     /* EState to use for "simple" expression evaluation */
+     EState       *simple_eval_estate;
+
      /* 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,
--- 947,954 ----
   * ----------
   */
  extern Datum plpgsql_exec_function(PLpgSQL_function *func,
!                       FunctionCallInfo fcinfo,
!                       EState *simple_eval_estate);
  extern HeapTuple plpgsql_exec_trigger(PLpgSQL_function *func,
                       TriggerData *trigdata);
  extern void plpgsql_exec_event_trigger(PLpgSQL_function *func,