From 7dc9589b5663abc31c00f9755c2943b3a5237bf6 Mon Sep 17 00:00:00 2001
From: Onder Kalaci <onder@citusdata.com>
Date: Wed, 28 Feb 2018 13:01:58 +0200
Subject: [PATCH] Handle failures during I/O

This commit checks the connection status right after any IO happens
on the socket.

This is necessary since before this commit we didn't pass any information
to the higher level functions whether we're done with the connection
(e.g., no IO required anymore) or an errors happened during the IO.
---
 .../distributed/connection/remote_commands.c       | 14 ++++++++++++++
 .../distributed/executor/multi_router_executor.c   | 12 ++++++++++++
 src/test/regress/expected/multi_extension.out      |  4 ++--
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/src/backend/distributed/connection/remote_commands.c b/src/backend/distributed/connection/remote_commands.c
index 2f6ef9de6..255541de8 100644
--- a/src/backend/distributed/connection/remote_commands.c
+++ b/src/backend/distributed/connection/remote_commands.c
@@ -106,6 +106,14 @@ ClearResults(MultiConnection *connection, bool raiseErrors)
 			MarkRemoteTransactionFailed(connection, raiseErrors);
 
 			success = false;
+
+			/* an error happened, there is nothing we can do more */
+			if (PQresultStatus(result) == PGRES_FATAL_ERROR)
+			{
+				PQclear(result);
+
+				break;
+			}
 		}
 
 		PQclear(result);
@@ -533,6 +541,12 @@ GetRemoteCommandResult(MultiConnection *connection, bool raiseInterrupts)
 
 	if (!FinishConnectionIO(connection, raiseInterrupts))
 	{
+		/* some error(s) happened while doing the I/O, signal the callers */
+		if (PQstatus(pgConn) == CONNECTION_BAD)
+		{
+			return PQmakeEmptyPGresult(pgConn, PGRES_FATAL_ERROR);
+		}
+
 		return NULL;
 	}
 
diff --git a/src/backend/distributed/executor/multi_router_executor.c b/src/backend/distributed/executor/multi_router_executor.c
index 2f8eeb9b4..3c14822f1 100644
--- a/src/backend/distributed/executor/multi_router_executor.c
+++ b/src/backend/distributed/executor/multi_router_executor.c
@@ -1446,6 +1446,12 @@ StoreQueryResult(CitusScanState *scanState, MultiConnection *connection,
 
 			commandFailed = true;
 
+			/* an error happened, there is nothing we can do more */
+			if (resultStatus == PGRES_FATAL_ERROR)
+			{
+				break;
+			}
+
 			/* continue, there could be other lingering results due to row mode */
 			continue;
 		}
@@ -1565,6 +1571,12 @@ ConsumeQueryResult(MultiConnection *connection, bool failOnError, int64 *rows)
 
 			commandFailed = true;
 
+			/* an error happened, there is nothing we can do more */
+			if (status == PGRES_FATAL_ERROR)
+			{
+				break;
+			}
+
 			/* continue, there could be other lingering results due to row mode */
 			continue;
 		}
diff --git a/src/test/regress/expected/multi_extension.out b/src/test/regress/expected/multi_extension.out
index 6272be6b7..19f9c9e63 100644
--- a/src/test/regress/expected/multi_extension.out
+++ b/src/test/regress/expected/multi_extension.out
@@ -369,8 +369,8 @@ SELECT create_distributed_table('t1', 'a');
 WARNING:  function assign_distributed_transaction_id(integer, integer, unknown) does not exist
 HINT:  No function matches the given name and argument types. You might need to add explicit type casts.
 CONTEXT:  while executing command on localhost:57637
-ERROR:  current transaction is aborted, commands ignored until end of transaction block
-CONTEXT:  while executing command on localhost:57637
+ERROR:  connection error: localhost:57637
+DETAIL:  another command is already in progress
 \c regression
 \c - - - :worker_1_port
 DROP DATABASE another;