Track PQresult allocations in server side (#1804)

When querying pg_locks (or using `pg_lock_status()`), approximately 75% of
backend memory allocations for resulting tuples weren't registered with
Vmtracker or Resource Group Control. This memory would also leak if the query
was cancelled or failed.

This happened because `CdbDispatchCommand()`, which is used by `pg_locks`,
internally uses libpq to obtain the results, which were allocated as PQresult
structures with bare `malloc()`, even on the server side.

This patch fixes both untracked memory issues by enforcing Vmtracker routines
for `PGresult` allocations on the server-side.

Please note that the original error from ADBDEV-7145, which is an
`AbortTransaction()` recurstion due to an OOM error while cancelling a query,
is fixed only in the next commit.

Ticket: ADBDEV-7691

Author: bandetto

src/backend/cdb/dispatcher/cdbdisp_async.c

@@ -685,6 +685,8 @@ checkDispatchResult(CdbDispatcherState *ds, int timeout_sec)
 			handlePollSuccess(pParms, fds);
 	}
 
+	SIMPLE_FAULT_INJECTOR("check_dispatch_result_end");
+
 	pfree(fds);
 }
 

src/interfaces/libpq/fe-exec.c

@@ -151,7 +151,7 @@ PQmakeEmptyPGresult(PGconn *conn, ExecStatusType status)
 {
 	PGresult   *result;
 
-	result = (PGresult *) malloc(sizeof(PGresult));
+	result = (PGresult *) pqPalloc(sizeof(PGresult));
 	if (!result)
 		return NULL;
 
@@ -407,7 +407,7 @@ dupEvents(PGEvent *events, int count)
 	if (!events || count <= 0)
 		return NULL;
 
-	newEvents = (PGEvent *) malloc(count * sizeof(PGEvent));
+	newEvents = (PGEvent *) pqPalloc(count * sizeof(PGEvent));
 	if (!newEvents)
 		return NULL;
 
@@ -417,12 +417,12 @@ dupEvents(PGEvent *events, int count)
 		newEvents[i].passThrough = events[i].passThrough;
 		newEvents[i].data = NULL;
 		newEvents[i].resultInitialized = FALSE;
-		newEvents[i].name = strdup(events[i].name);
+		newEvents[i].name = pqPstrdup(events[i].name);
 		if (!newEvents[i].name)
 		{
 			while (--i >= 0)
-				free(newEvents[i].name);
-			free(newEvents);
+				pqPfree(newEvents[i].name);
+			pqPfree(newEvents);
 			return NULL;
 		}
 	}
@@ -585,7 +585,7 @@ pqResultAlloc(PGresult *res, size_t nBytes, bool isBinary)
 	 */
 	if (nBytes >= PGRESULT_SEP_ALLOC_THRESHOLD)
 	{
-		block = (PGresult_data *) malloc(nBytes + PGRESULT_BLOCK_OVERHEAD);
+		block = (PGresult_data *) pqPalloc(nBytes + PGRESULT_BLOCK_OVERHEAD);
 		if (!block)
 			return NULL;
 		space = block->space + PGRESULT_BLOCK_OVERHEAD;
@@ -609,7 +609,7 @@ pqResultAlloc(PGresult *res, size_t nBytes, bool isBinary)
 	}
 
 	/* Otherwise, start a new block. */
-	block = (PGresult_data *) malloc(PGRESULT_DATA_BLOCKSIZE);
+	block = (PGresult_data *) pqPalloc(PGRESULT_DATA_BLOCKSIZE);
 	if (!block)
 		return NULL;
 	block->next = res->curBlock;
@@ -709,18 +709,18 @@ PQclear(PGresult *res)
 	}
 
 	if (res->events)
-		free(res->events);
+		pqPfree(res->events);
 
 	/* Free all the subsidiary blocks */
 	while ((block = res->curBlock) != NULL)
 	{
 		res->curBlock = block->next;
-		free(block);
+		pqPfree(block);
 	}
 
 	/* Free the top-level tuple pointer array */
 	if (res->tuples)
-		free(res->tuples);
+		pqPfree(res->tuples);
 
 	/* zero out the pointer fields to catch programming errors */
 	res->attDescs = NULL;
@@ -732,20 +732,20 @@ PQclear(PGresult *res)
 	/* res->curBlock was zeroed out earlier */
 
 	if (res->extras)
-		free(res->extras);
+		pqPfree(res->extras);
 	res->extraslen = 0;
 	res->extras = NULL;
 
 	if (res->aotupcounts)
-		free(res->aotupcounts);
+		pqPfree(res->aotupcounts);
 	res->naotupcounts = 0;
 
 	if (res->waitGxids)
-		free(res->waitGxids);
+		pqPfree(res->waitGxids);
 	res->waitGxids = NULL;
 	res->nWaits = 0;
 	/* Free the PGresult structure itself */
-	free(res);
+	pqPfree(res);
 }
 
 /*
@@ -951,10 +951,10 @@ pqAddTuple(PGresult *res, PGresAttValue *tup, const char **errmsgp)
 
 		if (res->tuples == NULL)
 			newTuples = (PGresAttValue **)
-				malloc(newSize * sizeof(PGresAttValue *));
+				pqPalloc(newSize * sizeof(PGresAttValue *));
 		else
 			newTuples = (PGresAttValue **)
-				realloc(res->tuples, newSize * sizeof(PGresAttValue *));
+				pqRepalloc(res->tuples, newSize * sizeof(PGresAttValue *));
 		if (!newTuples)
 			return FALSE;		/* malloc or realloc failed */
 		res->tupArrSize = newSize;

src/interfaces/libpq/fe-protocol3.c

@@ -267,7 +267,7 @@ pqParseInput3(PGconn *conn)
 
 					/* now just loop through */
 					conn->result->aotupcounts =
-						malloc(sizeof(PQaoRelTupCount) * conn->result->naotupcounts);
+						pqPalloc(sizeof(PQaoRelTupCount) * conn->result->naotupcounts);
 					ao = conn->result->aotupcounts;
 					for (i = 0; i < conn->result->naotupcounts; i++)
 					{
@@ -529,7 +529,7 @@ pqParseInput3(PGconn *conn)
 
 					if (pqGetInt(&conn->result->extraslen, 4, conn))
 						return;
-					conn->result->extras = malloc(conn->result->extraslen);
+					conn->result->extras = pqPalloc(conn->result->extraslen);
 					if (pqGetnchar((char *)conn->result->extras, conn->result->extraslen, conn))
 						return;
 					conn->asyncStatus = PGASYNC_READY;
@@ -555,7 +555,7 @@ pqParseInput3(PGconn *conn)
 					{
 						if (conn->result->waitGxids == NULL)
 							conn->result->waitGxids =
-								malloc(sizeof(int) * conn->result->nWaits);
+								pqPalloc(sizeof(int) * conn->result->nWaits);
 						for (i = 0; i < conn->result->nWaits; i++)
 						{
 							int gxid;

src/interfaces/libpq/libpq-int.h

@@ -83,6 +83,34 @@ typedef struct
 #endif
 #endif   /* USE_SSL */
 
+/*
+ * Definitions meant to lessen the malloc() usage for CDB routines in
+ * server-sided code for struct PGresult allocations.
+ */
+#ifndef FRONTEND
+#include "postgres.h"
+#include "utils/memutils.h"
+
+/*
+ * CurTransactionContext's lifetime lasts until the end of the current query,
+ * which does the job well for backends. Auxiliary processes do not set
+ * transaction contexts, so we have to use TopMemoryContext to achieve a
+ * lifetime equivalent to malloc().
+ */
+#define PQ_PALLOC_CONTEXT \
+	((CurTransactionContext != NULL) ? CurTransactionContext : TopMemoryContext)
+
+#define pqPalloc(sz) MemoryContextAlloc(PQ_PALLOC_CONTEXT, sz)
+#define pqPstrdup(x) MemoryContextStrdup(PQ_PALLOC_CONTEXT, x)
+#define pqRepalloc(x, sz) repalloc(x, sz)
+#define pqPfree(x) pfree(x)
+#else
+#define pqPalloc(sz) malloc(sz)
+#define pqPstrdup(x) strdup(x)
+#define pqRepalloc(x, sz) realloc(x, sz)
+#define pqPfree(x) free(x)
+#endif
+
 /*
  * POSTGRES backend dependent Constants.
  */

src/test/isolation2/input/resgroup/resgroup_oom.source

@@ -0,0 +1,49 @@
+-- start_matchsubs
+-- m/ERROR:  Out of memory.*/
+-- s/ERROR:  Out of memory.*/ERROR:  Out of memory/
+-- m/ERROR:  Canceling query because of high VMEM usage.*/
+-- s/ERROR:  Canceling query because of high VMEM usage.*/ERROR:  Out of memory/
+-- end_matchsubs
+
+-- start_ignore
+CREATE EXTENSION gp_inject_fault;
+DROP FUNCTION gp_mock_cdbdispatchcommand(amount int);
+DROP ROLE role_oom_test;
+DROP RESOURCE GROUP rg_oom_test;
+-- end_ignore
+CREATE RESOURCE GROUP rg_oom_test
+WITH (cpu_rate_limit=20, memory_limit=20, memory_shared_quota=100);
+CREATE ROLE role_oom_test RESOURCE GROUP rg_oom_test;
+CREATE OR REPLACE FUNCTION gp_mock_cdbdispatchcommand(amount int)
+RETURNS SETOF text AS '@abs_srcdir@/../regress/regress.so',
+'gp_mock_cdbdispatchcommand' LANGUAGE C;
+
+1: SET ROLE TO role_oom_test;
+
+-- Freeze coordinator's session after it reads results from segments.
+SELECT gp_inject_fault('check_dispatch_result_end', 'suspend', dbid)
+FROM gp_segment_configuration WHERE role = 'p' AND content = -1;
+
+-- Send the heavy query.
+1&: SELECT count(*) FROM gp_mock_cdbdispatchcommand(1000000);
+
+-- Wait until we receive everything.
+SELECT gp_wait_until_triggered_fault('check_dispatch_result_end', 1, dbid)
+FROM gp_segment_configuration WHERE role = 'p' AND content = -1;
+
+-- The query should've used ~135 MB of memory. Allow 15 MB error.
+WITH r AS (
+  SELECT (memory_usage->'-1'->'used')::text::int AS mb
+  FROM gp_toolkit.gp_resgroup_status WHERE rsgname = 'rg_oom_test'
+)
+SELECT r.mb < 150 AS "under 150 MB", r.mb > 120 AS "above 120 MB"
+FROM r;
+
+SELECT gp_inject_fault('check_dispatch_result_end', 'reset', dbid) FROM
+gp_segment_configuration WHERE role = 'p' AND content = -1;
+
+1<:
+
+DROP FUNCTION gp_mock_cdbdispatchcommand(amount int);
+DROP ROLE role_oom_test;
+DROP RESOURCE GROUP rg_oom_test;

src/test/isolation2/isolation2_resgroup_schedule

@@ -61,4 +61,6 @@ test: resgroup/resgroup_large_group_id
 
 test: resgroup/resgroup_startup_memory
 
+test: resgroup/resgroup_oom
+
 test: resgroup/disable_resgroup

src/test/isolation2/output/resgroup/resgroup_oom.source

@@ -0,0 +1,59 @@
+-- start_matchsubs
+-- m/ERROR:  Out of memory.*/
+-- s/ERROR:  Out of memory.*/ERROR:  Out of memory/
+-- m/ERROR:  Canceling query because of high VMEM usage.*/
+-- s/ERROR:  Canceling query because of high VMEM usage.*/ERROR:  Out of memory/
+-- end_matchsubs
+
+CREATE RESOURCE GROUP rg_oom_test WITH (cpu_rate_limit=20, memory_limit=20, memory_shared_quota=100);
+CREATE
+CREATE ROLE role_oom_test RESOURCE GROUP rg_oom_test;
+CREATE
+CREATE OR REPLACE FUNCTION gp_mock_cdbdispatchcommand(amount int) RETURNS SETOF text AS '@abs_srcdir@/../regress/regress.so', 'gp_mock_cdbdispatchcommand' LANGUAGE C;
+CREATE
+
+1: SET ROLE TO role_oom_test;
+SET
+
+-- Freeze coordinator's session after it reads results from segments.
+SELECT gp_inject_fault('check_dispatch_result_end', 'suspend', dbid) FROM gp_segment_configuration WHERE role = 'p' AND content = -1;
+ gp_inject_fault 
+-----------------
+ Success:        
+(1 row)
+
+-- Send the heavy query.
+1&: SELECT count(*) FROM gp_mock_cdbdispatchcommand(1000000);  <waiting ...>
+
+-- Wait until we receive everything.
+SELECT gp_wait_until_triggered_fault('check_dispatch_result_end', 1, dbid) FROM gp_segment_configuration WHERE role = 'p' AND content = -1;
+ gp_wait_until_triggered_fault 
+-------------------------------
+ Success:                      
+(1 row)
+
+-- The query should've used ~135 MB of memory. Allow 15 MB error.
+WITH r AS ( SELECT (memory_usage->'-1'->'used')::text::int AS mb FROM gp_toolkit.gp_resgroup_status WHERE rsgname = 'rg_oom_test' ) SELECT r.mb < 150 AS "under 150 MB", r.mb > 120 AS "above 120 MB" FROM r;
+ under 150 MB | above 120 MB 
+--------------+--------------
+ t            | t            
+(1 row)
+
+SELECT gp_inject_fault('check_dispatch_result_end', 'reset', dbid) FROM gp_segment_configuration WHERE role = 'p' AND content = -1;
+ gp_inject_fault 
+-----------------
+ Success:        
+(1 row)
+
+1<:  <... completed>
+ count   
+---------
+ 4000000 
+(1 row)
+
+DROP FUNCTION gp_mock_cdbdispatchcommand(amount int);
+DROP
+DROP ROLE role_oom_test;
+DROP
+DROP RESOURCE GROUP rg_oom_test;
+DROP