Modify pg_stat_get_activity to build a tuplestore

This updates pg_stat_get_activity() to build a tuplestore for its
results instead of using the old-style multiple-call method.  This
simplifies the function, though that wasn't the primary motivation for
the change, which is that we may turn it into a helper function which
can filter the results (or not) much more easily.

diff --git a/src/backend/utils/adt/pgstatfuncs.c b/src/backend/utils/adt/pgstatfuncs.c
index bbe94c34a15976ea76bb54a9673edfd928c5333d..2b3778b03ad85df70d725fdae6c272429042dd0b 100644 (file)
--- a/src/backend/utils/adt/pgstatfuncs.c
+++ b/src/backend/utils/adt/pgstatfuncs.c
@@ -524,137 +524,75 @@ pg_stat_get_backend_idset(PG_FUNCTION_ARGS)
        }
 }
 
+/*
+ * Returns activity of PG backends.
+ */
 Datum
 pg_stat_get_activity(PG_FUNCTION_ARGS)
 {
-       FuncCallContext *funcctx;
-
-       if (SRF_IS_FIRSTCALL())
-       {
-               MemoryContext oldcontext;
-               TupleDesc       tupdesc;
-
-               funcctx = SRF_FIRSTCALL_INIT();
-
-               oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
-
-               tupdesc = CreateTemplateTupleDesc(22, false);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 1, "datid",
-                                                  OIDOID, -1, 0);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 2, "pid",
-                                                  INT4OID, -1, 0);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 3, "usesysid",
-                                                  OIDOID, -1, 0);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 4, "application_name",
-                                                  TEXTOID, -1, 0);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 5, "state",
-                                                  TEXTOID, -1, 0);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 6, "query",
-                                                  TEXTOID, -1, 0);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 7, "waiting",
-                                                  BOOLOID, -1, 0);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 8, "act_start",
-                                                  TIMESTAMPTZOID, -1, 0);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 9, "query_start",
-                                                  TIMESTAMPTZOID, -1, 0);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 10, "backend_start",
-                                                  TIMESTAMPTZOID, -1, 0);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 11, "state_change",
-                                                  TIMESTAMPTZOID, -1, 0);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 12, "client_addr",
-                                                  INETOID, -1, 0);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 13, "client_hostname",
-                                                  TEXTOID, -1, 0);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 14, "client_port",
-                                                  INT4OID, -1, 0);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 15, "backend_xid",
-                                                  XIDOID, -1, 0);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 16, "backend_xmin",
-                                                  XIDOID, -1, 0);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 17, "ssl",
-                                                  BOOLOID, -1, 0);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 18, "sslversion",
-                                                  TEXTOID, -1, 0);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 19, "sslcipher",
-                                                  TEXTOID, -1, 0);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 20, "sslbits",
-                                                  INT4OID, -1, 0);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 21, "sslcompression",
-                                                  BOOLOID, -1, 0);
-               TupleDescInitEntry(tupdesc, (AttrNumber) 22, "sslclientdn",
-                                                  TEXTOID, -1, 0);
-
-               funcctx->tuple_desc = BlessTupleDesc(tupdesc);
-
-               funcctx->user_fctx = palloc0(sizeof(int));
-               if (PG_ARGISNULL(0))
-               {
-                       /* Get all backends */
-                       funcctx->max_calls = pgstat_fetch_stat_numbackends();
-               }
-               else
-               {
-                       /*
-                        * Get one backend - locate by pid.
-                        *
-                        * We lookup the backend early, so we can return zero rows if it
-                        * doesn't exist, instead of returning a single row full of NULLs.
-                        */
-                       int                     pid = PG_GETARG_INT32(0);
-                       int                     i;
-                       int                     n = pgstat_fetch_stat_numbackends();
-
-                       for (i = 1; i <= n; i++)
-                       {
-                               PgBackendStatus *be = pgstat_fetch_stat_beentry(i);
-
-                               if (be)
-                               {
-                                       if (be->st_procpid == pid)
-                                       {
-                                               *(int *) (funcctx->user_fctx) = i;
-                                               break;
-                                       }
-                               }
-                       }
-
-                       if (*(int *) (funcctx->user_fctx) == 0)
-                               /* Pid not found, return zero rows */
-                               funcctx->max_calls = 0;
-                       else
-                               funcctx->max_calls = 1;
-               }
-
-               MemoryContextSwitchTo(oldcontext);
-       }
-
-       /* stuff done on every call of the function */
-       funcctx = SRF_PERCALL_SETUP();
-
-       if (funcctx->call_cntr < funcctx->max_calls)
+#define PG_STAT_GET_ACTIVITY_COLS      22
+       int                                     num_backends = pgstat_fetch_stat_numbackends();
+       int                                     curr_backend;
+       int                                     pid = PG_ARGISNULL(0) ? -1 : PG_GETARG_INT32(0);
+       ReturnSetInfo      *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
+       TupleDesc                       tupdesc;
+       Tuplestorestate    *tupstore;
+       MemoryContext           per_query_ctx;
+       MemoryContext           oldcontext;
+
+       /* check to see if caller supports us returning a tuplestore */
+       if (rsinfo == NULL || !IsA(rsinfo, ReturnSetInfo))
+               ereport(ERROR,
+                               (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+                                errmsg("set-valued function called in context that cannot accept a set")));
+       if (!(rsinfo->allowedModes & SFRM_Materialize))
+               ereport(ERROR,
+                               (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+                                errmsg("materialize mode required, but it is not " \
+                                               "allowed in this context")));
+
+       /* Build a tuple descriptor for our result type */
+       if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
+               elog(ERROR, "return type must be a row type");
+
+       per_query_ctx = rsinfo->econtext->ecxt_per_query_memory;
+       oldcontext = MemoryContextSwitchTo(per_query_ctx);
+
+       tupstore = tuplestore_begin_heap(true, false, work_mem);
+       rsinfo->returnMode = SFRM_Materialize;
+       rsinfo->setResult = tupstore;
+       rsinfo->setDesc = tupdesc;
+
+       MemoryContextSwitchTo(oldcontext);
+
+       /* 1-based index */
+       for (curr_backend = 1; curr_backend <= num_backends; curr_backend++)
        {
                /* for each row */
-               Datum           values[22];
-               bool            nulls[22];
-               HeapTuple       tuple;
+               Datum           values[PG_STAT_GET_ACTIVITY_COLS];
+               bool            nulls[PG_STAT_GET_ACTIVITY_COLS];
                LocalPgBackendStatus *local_beentry;
                PgBackendStatus *beentry;
 
                MemSet(values, 0, sizeof(values));
                MemSet(nulls, 0, sizeof(nulls));
 
-               if (*(int *) (funcctx->user_fctx) > 0)
-               {
-                       /* Get specific pid slot */
-                       local_beentry = pgstat_fetch_stat_local_beentry(*(int *) (funcctx->user_fctx));
-                       beentry = &local_beentry->backendStatus;
-               }
-               else
+               if (pid != -1)
                {
-                       /* Get the next one in the list */
-                       local_beentry = pgstat_fetch_stat_local_beentry(funcctx->call_cntr + 1);        /* 1-based index */
-                       beentry = &local_beentry->backendStatus;
+                       /* Skip any which are not the one we're looking for. */
+                       PgBackendStatus *be = pgstat_fetch_stat_beentry(curr_backend);
+
+                       if (!be || be->st_procpid != pid)
+                               continue;
+
                }
+
+               /* Get the next one in the list */
+               local_beentry = pgstat_fetch_stat_local_beentry(curr_backend);
+               if (!local_beentry)
+                       continue;
+
+               beentry = &local_beentry->backendStatus;
                if (!beentry)
                {
                        int                     i;
@@ -665,8 +603,8 @@ pg_stat_get_activity(PG_FUNCTION_ARGS)
                        nulls[5] = false;
                        values[5] = CStringGetTextDatum("<backend information not available>");
 
-                       tuple = heap_form_tuple(funcctx->tuple_desc, values, nulls);
-                       SRF_RETURN_NEXT(funcctx, HeapTupleGetDatum(tuple));
+                       tuplestore_putvalues(tupstore, tupdesc, values, nulls);
+                       continue;
                }
 
                /* Values available to all callers */
@@ -839,15 +777,17 @@ pg_stat_get_activity(PG_FUNCTION_ARGS)
                        nulls[13] = true;
                }
 
-               tuple = heap_form_tuple(funcctx->tuple_desc, values, nulls);
+               tuplestore_putvalues(tupstore, tupdesc, values, nulls);
 
-               SRF_RETURN_NEXT(funcctx, HeapTupleGetDatum(tuple));
-       }
-       else
-       {
-               /* nothing left */
-               SRF_RETURN_DONE(funcctx);
+               /* If only a single backend was requested, and we found it, break. */
+               if (pid != -1)
+                       break;
        }
+
+       /* clean up and return the tuplestore */
+       tuplestore_donestoring(tupstore);
+
+       return (Datum) 0;
 }