From 095b5a3cd0a482623b3eaf9bf83c9002e4c32fd6 Mon Sep 17 00:00:00 2001
From: Alexander Kuzmenkov <akuzmenkov@timescale.com>
Date: Mon, 5 Sep 2022 13:51:42 +0300
Subject: [PATCH] Fix a possible error while flushing the COPY data

Apparently, pqFlush can report a success status, but the underlying
socket can become invalid at the same time. Check for this.

In passing, also change to the proper error reporting functions that
use the connection error information, where needed.
---
 tsl/src/remote/connection.c | 43 ++++++++++++++++++++++++++-----------
 1 file changed, 30 insertions(+), 13 deletions(-)

diff --git a/tsl/src/remote/connection.c b/tsl/src/remote/connection.c
index 0235ead4978..025b2f98fb6 100644
--- a/tsl/src/remote/connection.c
+++ b/tsl/src/remote/connection.c
@@ -2255,13 +2255,30 @@ remote_connection_end_copy(TSConnection *conn, TSConnectionError *err)
 
 			if (flush_result == 1)
 			{
+				/*
+				 * In some rare cases, flush might report that it's busy, but
+				 * actually there was an error and the socket became invalid.
+				 * Check for it. This is something we have observed in COPY
+				 * queries used for performance testing with tsbench, but not
+				 * sure how it happens exactly, must be in the depths of
+				 * pqReadData called by pqFlush.
+				 */
+				int socket = PQsocket(conn->pg_conn);
+				if (socket == PGINVALID_SOCKET)
+				{
+					return fill_connection_error(err,
+												 ERRCODE_CONNECTION_EXCEPTION,
+												 "failed to flush the COPY connection",
+												 conn);
+				}
+
 				/*
 				 * The socket is busy, wait. We don't care about the wait result
 				 * here, because whether it is a timeout or the socket became
 				 * writeable, we just retry.
 				 */
 				(void) WaitLatchOrSocket(MyLatch,
-										 PQsocket(conn->pg_conn),
+										 socket,
 										 WL_TIMEOUT | WL_SOCKET_WRITEABLE,
 										 /* timeout = */ 1000,
 										 /* wait_event_info = */ 0);
@@ -2274,20 +2291,20 @@ remote_connection_end_copy(TSConnection *conn, TSConnectionError *err)
 			else
 			{
 				/* Error. */
-				return fill_simple_error(err,
-										 ERRCODE_CONNECTION_EXCEPTION,
-										 "failed to flush the COPY connection",
-										 conn);
+				return fill_connection_error(err,
+											 ERRCODE_CONNECTION_EXCEPTION,
+											 "failed to flush the COPY connection",
+											 conn);
 			}
 		}
 
 		/* Switch the connection into blocking mode. */
 		if (PQsetnonblocking(conn->pg_conn, 0) != 0)
 		{
-			return fill_simple_error(err,
-									 ERRCODE_CONNECTION_EXCEPTION,
-									 "failed to set the connection into blocking mode",
-									 conn);
+			return fill_connection_error(err,
+										 ERRCODE_CONNECTION_EXCEPTION,
+										 "failed to set the connection into blocking mode",
+										 conn);
 		}
 	}
 
@@ -2318,10 +2335,10 @@ remote_connection_end_copy(TSConnection *conn, TSConnectionError *err)
 		return false;
 
 	if (PQputCopyEnd(conn->pg_conn, NULL) != 1)
-		return fill_simple_error(err,
-								 ERRCODE_CONNECTION_EXCEPTION,
-								 "could not end remote COPY",
-								 conn);
+		return fill_connection_error(err,
+									 ERRCODE_CONNECTION_EXCEPTION,
+									 "could not end remote COPY",
+									 conn);
 
 	success = true;
 	remote_connection_set_status(conn, CONN_PROCESSING);