1 files changed, 54 insertions, 37 deletions

diff --git a/src/interfaces/libpq/fe-exec.c b/src/interfaces/libpq/fe-exec.c
index ac969e7b3f8..62026538264 100644
--- a/src/interfaces/libpq/fe-exec.c
+++ b/src/interfaces/libpq/fe-exec.c
@@ -791,6 +791,32 @@ pqSaveErrorResult(PGconn *conn)
 }
 
 /*
+ * As above, and append conn->write_err_msg to whatever other error we have.
+ * This is used when we've detected a write failure and have exhausted our
+ * chances of reporting something else instead.
+ */
+static void
+pqSaveWriteError(PGconn *conn)
+{
+	/*
+	 * Ensure conn->result is an error result, and add anything in
+	 * conn->errorMessage to it.
+	 */
+	pqSaveErrorResult(conn);
+
+	/*
+	 * Now append write_err_msg to that.  If it's null because of previous
+	 * strdup failure, do what we can.  (It's likely our machinations here are
+	 * all getting OOM failures as well, but ...)
+	 */
+	if (conn->write_err_msg && conn->write_err_msg[0] != '\0')
+		pqCatenateResultError(conn->result, conn->write_err_msg);
+	else
+		pqCatenateResultError(conn->result,
+							  libpq_gettext("write to server failed\n"));
+}
+
+/*
  * This subroutine prepares an async result object for return to the caller.
  * If there is not already an async result object, build an error object
  * using whatever is in conn->errorMessage.  In any case, clear the async
@@ -1224,7 +1250,7 @@ PQsendQuery(PGconn *conn, const char *query)
 		pqPuts(query, conn) < 0 ||
 		pqPutMsgEnd(conn) < 0)
 	{
-		pqHandleSendFailure(conn);
+		/* error message should be set up already */
 		return 0;
 	}
 
@@ -1243,7 +1269,7 @@ PQsendQuery(PGconn *conn, const char *query)
 	 */
 	if (pqFlush(conn) < 0)
 	{
-		pqHandleSendFailure(conn);
+		/* error message should be set up already */
 		return 0;
 	}
 
@@ -1389,7 +1415,7 @@ PQsendPrepare(PGconn *conn,
 	return 1;
 
 sendFailed:
-	pqHandleSendFailure(conn);
+	/* error message should be set up already */
 	return 0;
 }
 
@@ -1641,40 +1667,11 @@ PQsendQueryGuts(PGconn *conn,
 	return 1;
 
 sendFailed:
-	pqHandleSendFailure(conn);
+	/* error message should be set up already */
 	return 0;
 }
 
 /*
- * pqHandleSendFailure: try to clean up after failure to send command.
- *
- * Primarily, what we want to accomplish here is to process any ERROR or
- * NOTICE messages that the backend might have sent just before it died.
- * Since we're in IDLE state, all such messages will get sent to the notice
- * processor.
- *
- * NOTE: this routine should only be called in PGASYNC_IDLE state.
- */
-void
-pqHandleSendFailure(PGconn *conn)
-{
-	/*
-	 * Accept and parse any available input data, ignoring I/O errors.  Note
-	 * that if pqReadData decides the backend has closed the channel, it will
-	 * close our side of the socket --- that's just what we want here.
-	 */
-	while (pqReadData(conn) > 0)
-		parseInput(conn);
-
-	/*
-	 * Be sure to parse available input messages even if we read no data.
-	 * (Note: calling parseInput within the above loop isn't really necessary,
-	 * but it prevents buffer bloat if there's a lot of data available.)
-	 */
-	parseInput(conn);
-}
-
-/*
  * Select row-by-row processing mode
  */
 int
@@ -1763,8 +1760,11 @@ PQisBusy(PGconn *conn)
 	/* Parse any available data, if our state permits. */
 	parseInput(conn);
 
-	/* PQgetResult will return immediately in all states except BUSY. */
-	return conn->asyncStatus == PGASYNC_BUSY;
+	/*
+	 * PQgetResult will return immediately in all states except BUSY, or if we
+	 * had a write failure.
+	 */
+	return conn->asyncStatus == PGASYNC_BUSY || conn->write_failed;
 }
 
 
@@ -1804,7 +1804,13 @@ PQgetResult(PGconn *conn)
 			}
 		}
 
-		/* Wait for some more data, and load it. */
+		/*
+		 * Wait for some more data, and load it.  (Note: if the connection has
+		 * been lost, pqWait should return immediately because the socket
+		 * should be read-ready, either with the last server data or with an
+		 * EOF indication.  We expect therefore that this won't result in any
+		 * undue delay in reporting a previous write failure.)
+		 */
 		if (flushResult ||
 			pqWait(true, false, conn) ||
 			pqReadData(conn) < 0)
@@ -1820,6 +1826,17 @@ PQgetResult(PGconn *conn)
 
 		/* Parse it. */
 		parseInput(conn);
+
+		/*
+		 * If we had a write error, but nothing above obtained a query result
+		 * or detected a read error, report the write error.
+		 */
+		if (conn->write_failed && conn->asyncStatus == PGASYNC_BUSY)
+		{
+			pqSaveWriteError(conn);
+			conn->asyncStatus = PGASYNC_IDLE;
+			return pqPrepareAsyncResult(conn);
+		}
 	}
 
 	/* Return the appropriate thing. */
@@ -2252,7 +2269,7 @@ PQsendDescribe(PGconn *conn, char desc_type, const char *desc_target)
 	return 1;
 
 sendFailed:
-	pqHandleSendFailure(conn);
+	/* error message should be set up already */
 	return 0;
 }