Gracefully handle OOM when cancelling a query (#1804)

Previosly, if a query fails under resource group manager due to a lack of
memory, the query is cancelled due to an error.

The coordinator would still try to receive tuples from segments upon
cancellation. Retrieving tuples requires more memory allocations, which may
cause the same OOM error again. The query will get cancelled due to an error,
entering into a recursion and leading to an eventual crash.

This patch implements proper query cancellation by explicitly ignoring any
libpq data from segments without extra memory allocations.

Ticket: ADBDEV-7690

Author: bandetto

src/backend/cdb/dispatcher/cdbdisp_async.c

@@ -135,6 +135,9 @@ static bool processResults(CdbDispatchResult *dispatchResult);
 static void
 			signalQEs(CdbDispatchCmdAsync *pParms);
 
+static void
+			signalQE(CdbDispatchResult *dispatchResult, DispatchWaitMode waitMode);
+
 static void
 			checkSegmentAlive(CdbDispatchCmdAsync *pParms);
 
@@ -685,6 +688,8 @@ checkDispatchResult(CdbDispatcherState *ds, int timeout_sec)
 			handlePollSuccess(pParms, fds);
 	}
 
+	SIMPLE_FAULT_INJECTOR("check_dispatch_result_end");
+
 	pfree(fds);
 }
 
@@ -820,6 +825,60 @@ handlePollError(CdbDispatchCmdAsync *pParms)
 	return;
 }
 
+/*
+ * Try to completely discard results from a dispatchResult's connection without
+ * extra memory allocations by abusing libpq state-machine hacks.
+ */
+static void
+tryToWipeResults(CdbDispatchResult *dispatchResult)
+{
+	PGresult   *res;
+	PGnotify   *notify;
+	PGconn	   *conn = dispatchResult->segdbDesc->conn;
+
+	/* Free some memory and replace current result with a fatal error dummy. */
+	pqSaveErrorResult(conn);
+
+	/* Make sure PQgetResult() calls are not blocking. */
+	PQconsumeInput(conn);
+
+	/*
+	 * Discard anything that is unread. Since our result contains a fatal
+	 * error, we'll just consume the entire message without actually parsing
+	 * it.
+	 */
+	while (!PQisBusy(conn) && (res = PQgetResult(conn)) != NULL)
+	{
+		switch (PQresultStatus(res))
+		{
+			case PGRES_COPY_IN:
+			case PGRES_COPY_OUT:
+			case PGRES_COPY_BOTH:
+				PQendcopy(conn);
+				/* fallthrough */
+			default:
+				PQclear(res);
+		}
+
+		pqSaveErrorResult(conn);
+	}
+
+	/* Free notices. */
+	while ((notify = PQnotifies(conn)) != NULL)
+		PQfreemem(notify);
+
+	/* The result is not needed anymore. */
+	pqClearAsyncResult(conn);
+
+	if (PQisBusy(conn) && PQstatus(conn) != CONNECTION_BAD)
+	{
+		/* Some work is still remaining until we can die. */
+		return;
+	}
+
+	dispatchResult->stillRunning = false;
+}
+
 /*
  * Receive and process results from QEs.
  */
@@ -866,6 +925,24 @@ handlePollSuccess(CdbDispatchCmdAsync *pParms,
 		ELOG_DISPATCHER_DEBUG("PQsocket says there are results from %d of %d (%s)",
 							  i + 1, pParms->dispatchCount, segdbDesc->whoami);
 
+
+		/*
+		 * Was dispatchCancel() the callee? We don't need to read the results then.
+		 */
+		if (pParms->waitMode == DISPATCH_WAIT_CANCEL)
+		{
+			/*
+			 * If we're cancelling the transaction due to an OOM, there
+			 * might not be enough memory to discard the result properly.
+			 * Let's get the big guns out.
+			 */
+			tryToWipeResults(dispatchResult);
+
+			forwardQENotices();
+
+			continue;
+		}
+
 		/*
 		 * Receive and process results from this QE.
 		 */
@@ -912,38 +989,41 @@ static void
 signalQEs(CdbDispatchCmdAsync *pParms)
 {
 	int			i;
-	DispatchWaitMode waitMode = pParms->waitMode;
 
 	for (i = 0; i < pParms->dispatchCount; i++)
-	{
-		char		errbuf[256];
-		bool		sent = false;
-		CdbDispatchResult *dispatchResult = pParms->dispatchResultPtrArray[i];
+		signalQE(pParms->dispatchResultPtrArray[i], pParms->waitMode);
+}
 
-		Assert(dispatchResult != NULL);
-		SegmentDatabaseDescriptor *segdbDesc = dispatchResult->segdbDesc;
+/*
+ * Send finish or cancel signal to QE if needed.
+ */
+static void
+signalQE(CdbDispatchResult *dispatchResult, DispatchWaitMode waitMode)
+{
+	Assert(dispatchResult != NULL);
+	SegmentDatabaseDescriptor *segdbDesc = dispatchResult->segdbDesc;
 
-		/*
-		 * Don't send the signal if - QE is finished or canceled - the signal
-		 * was already sent - connection is dead
-		 */
+	/*
+	 * Don't send the signal if - QE is finished or canceled - the signal
+	 * was already sent - connection is dead
+	 */
 
-		if (!dispatchResult->stillRunning ||
-			dispatchResult->wasCanceled ||
-			(pParms->waitMode == DISPATCH_WAIT_ACK_ROOT &&
-			 dispatchResult->receivedAckMsg) ||
-			cdbconn_isBadConnection(segdbDesc))
-			continue;
+	if (!dispatchResult->stillRunning ||
+		dispatchResult->wasCanceled ||
+		(waitMode == DISPATCH_WAIT_ACK_ROOT &&
+		 dispatchResult->receivedAckMsg) ||
+		cdbconn_isBadConnection(segdbDesc))
+	{
+		return;
+	}
 
-		memset(errbuf, 0, sizeof(errbuf));
+	char		errbuf[256] = {0};
 
-		sent = cdbconn_signalQE(segdbDesc, errbuf, waitMode == DISPATCH_WAIT_CANCEL);
-		if (sent)
-			dispatchResult->sentSignal = waitMode;
-		else
-			elog(LOG, "Unable to cancel: %s",
-				 strlen(errbuf) == 0 ? "cannot allocate PGCancel" : errbuf);
-	}
+	if (cdbconn_signalQE(segdbDesc, errbuf, waitMode == DISPATCH_WAIT_CANCEL))
+		dispatchResult->sentSignal = waitMode;
+	else
+		elog(LOG, "Unable to cancel: %s",
+			 strlen(errbuf) == 0 ? "cannot allocate PGCancel" : errbuf);
 }
 
 /*

src/test/isolation2/input/parallel_retrieve_cursor/fault_inject.source

@@ -249,3 +249,25 @@ SELECT gp_inject_fault('fetch_tuples_from_endpoint', 'resume', dbid)
     WHERE content=2 AND role='p';
 
 DROP TABLE t2;
+
+SELECT gp_inject_fault('alloc_endpoint_slot_full', 'reset', 2);
+SELECT gp_inject_fault('alloc_endpoint_slot_full_reset', 'reset', 2);
+
+-- Test 8: QD shouldn't hang waiting.
+
+DROP TABLE IF EXISTS t3;
+CREATE TABLE t3 AS SELECT generate_series(1, 10) AS a DISTRIBUTED by (a);
+
+SELECT gp_inject_fault('alloc_endpoint_slot_full', 'error', dbid)
+FROM gp_segment_configuration WHERE role = 'p' AND content = 0;
+
+BEGIN;
+-- QE encounters error and destroys the endpoint immediately. QD successfully
+-- resets the query, sends ACKs, and discards transaction.
+DECLARE c1 PARALLEL RETRIEVE CURSOR FOR SELECT * FROM t3;
+ROLLBACK;
+
+SELECT gp_inject_fault('alloc_endpoint_slot_full', 'reset', dbid)
+FROM gp_segment_configuration WHERE role = 'p' AND content = 0;
+
+DROP TABLE t3;

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

@@ -44,6 +44,10 @@ gp_segment_configuration WHERE role = 'p' AND content = -1;
 
 1<:
 
+-- And finally, make sure we don't enter error recursion on fail.
+ALTER RESOURCE GROUP rg_oom_test SET memory_shared_quota 0;
+1: SELECT count(*) FROM gp_mock_cdbdispatchcommand(10000000);
+
 DROP FUNCTION gp_mock_cdbdispatchcommand(amount int);
 DROP ROLE role_oom_test;
 DROP RESOURCE GROUP rg_oom_test;

src/test/isolation2/output/parallel_retrieve_cursor/fault_inject.source

@@ -770,3 +770,45 @@ ROLLBACK
 
 DROP TABLE t2;
 DROP
+
+SELECT gp_inject_fault('alloc_endpoint_slot_full', 'reset', 2);
+ gp_inject_fault 
+-----------------
+ Success:        
+(1 row)
+SELECT gp_inject_fault('alloc_endpoint_slot_full_reset', 'reset', 2);
+ gp_inject_fault 
+-----------------
+ Success:        
+(1 row)
+
+-- Test 8: QD shouldn't hang waiting.
+
+DROP TABLE IF EXISTS t3;
+DROP
+CREATE TABLE t3 AS SELECT generate_series(1, 10) AS a DISTRIBUTED by (a);
+CREATE 10
+
+SELECT gp_inject_fault('alloc_endpoint_slot_full', 'error', dbid) FROM gp_segment_configuration WHERE role = 'p' AND content = 0;
+ gp_inject_fault 
+-----------------
+ Success:        
+(1 row)
+
+BEGIN;
+BEGIN
+-- QE encounters error and destroys the endpoint immediately. QD successfully
+-- resets the query, sends ACKs, and discards transaction.
+DECLARE c1 PARALLEL RETRIEVE CURSOR FOR SELECT * FROM t3;
+ERROR:  fault triggered, fault name:'alloc_endpoint_slot_full' fault type:'error'  (seg0 127.0.0.1:6002 pid=8916)
+ROLLBACK;
+ROLLBACK
+
+SELECT gp_inject_fault('alloc_endpoint_slot_full', 'reset', dbid) FROM gp_segment_configuration WHERE role = 'p' AND content = 0;
+ gp_inject_fault 
+-----------------
+ Success:        
+(1 row)
+
+DROP TABLE t3;
+DROP

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

@@ -51,6 +51,13 @@ SELECT gp_inject_fault('check_dispatch_result_end', 'reset', dbid) FROM gp_segme
  4000000 
 (1 row)
 
+-- And finally, make sure we don't enter error recursion on fail.
+ALTER RESOURCE GROUP rg_oom_test SET memory_shared_quota 0;
+ALTER
+1: SELECT count(*) FROM gp_mock_cdbdispatchcommand(10000000);
+ERROR:  Out of memory
+DETAIL:  Resource group memory limit reached
+
 DROP FUNCTION gp_mock_cdbdispatchcommand(amount int);
 DROP
 DROP ROLE role_oom_test;

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

@@ -0,0 +1,65 @@
+-- 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)
+
+-- And finally, make sure we don't enter error recursion on fail.
+ALTER RESOURCE GROUP rg_oom_test SET memory_shared_quota 0;
+ALTER
+1: SELECT count(*) FROM gp_mock_cdbdispatchcommand(10000000);
+ERROR:  Out of memory
+
+DROP FUNCTION gp_mock_cdbdispatchcommand(amount int);
+DROP
+DROP ROLE role_oom_test;
+DROP
+DROP RESOURCE GROUP rg_oom_test;
+DROP