pg_stat_get_backend_idset should reset fmgr_info->fn_extra at end of

execution, so that it restarts correctly if query tree is used again.

diff --git a/src/backend/utils/adt/pgstatfuncs.c b/src/backend/utils/adt/pgstatfuncs.c
index 23617b850753a86cfa1430b2f456d8d4952ff951..01cc5e45b399ab17d9aaccf27cd834f95836406f 100644 (file)
--- a/src/backend/utils/adt/pgstatfuncs.c
+++ b/src/backend/utils/adt/pgstatfuncs.c
@@ -179,33 +179,29 @@ Datum
 pg_stat_get_backend_idset(PG_FUNCTION_ARGS)
 {
    FmgrInfo       *fmgr_info = fcinfo->flinfo;
-   MemoryContext   oldcxt;
    int32           result;
 
-   if (fcinfo->resultinfo == NULL)
-       elog(ERROR, "called in context that does not accept a set result (1)");
-   if (!IsA(fcinfo->resultinfo, ReturnSetInfo))
-       elog(ERROR, "called in context that does not accept a set result (2)");
+   if (fcinfo->resultinfo == NULL ||
+       !IsA(fcinfo->resultinfo, ReturnSetInfo))
+       elog(ERROR, "pg_stat_get_backend_idset: called in context that does not accept a set result");
 
    if (fmgr_info->fn_extra == NULL)
    {
        if (fmgr_info->fn_mcxt == NULL)
            elog(ERROR, "No function memory context in set-function");
-
-       oldcxt = MemoryContextSwitchTo(fmgr_info->fn_mcxt);
-       fmgr_info->fn_extra = palloc(sizeof(int32) * 2);
-       ((int32 *)(fmgr_info->fn_extra))[0] = 0;
-       ((int32 *)(fmgr_info->fn_extra))[1] = 
-                       (int32)pgstat_fetch_stat_numbackends();
-       MemoryContextSwitchTo(oldcxt);
+       fmgr_info->fn_extra = MemoryContextAlloc(fmgr_info->fn_mcxt,
+                                                2 * sizeof(int));
+       ((int *)(fmgr_info->fn_extra))[0] = 0;
+       ((int *)(fmgr_info->fn_extra))[1] = pgstat_fetch_stat_numbackends();
    }
 
-   ((int32 *)(fmgr_info->fn_extra))[0] += 1;
-   result = ((int32 *)(fmgr_info->fn_extra))[0];
-
+   ((int *)(fmgr_info->fn_extra))[0] += 1;
+   result = ((int *)(fmgr_info->fn_extra))[0];
 
-   if (result > ((int32 *)(fmgr_info->fn_extra))[1])
+   if (result > ((int *)(fmgr_info->fn_extra))[1])
    {
+       pfree(fmgr_info->fn_extra);
+       fmgr_info->fn_extra = NULL;
        ((ReturnSetInfo *)(fcinfo->resultinfo))->isDone = ExprEndResult;
        PG_RETURN_NULL();
    }