From 088a743d39efc26a051ee2529a461ffdaa3c5efc Mon Sep 17 00:00:00 2001
From: Matthias van de Meent <matthias@neon.tech>
Date: Thu, 27 Jun 2024 10:58:15 +0200
Subject: [PATCH] Handle PS error responses cleanly

Before this, we'd assume the PS connection would be good once we got a response.
Now, we also check that the response was a positive one.

This means our exponential backoff policy now also applies to PS connections that
do connect but fail to get a pagestream, where previously it did not.
---
 pgxn/neon/libpagestore.c | 43 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 43 insertions(+)

diff --git a/pgxn/neon/libpagestore.c b/pgxn/neon/libpagestore.c
index a665cafafe..9327e5f00d 100644
--- a/pgxn/neon/libpagestore.c
+++ b/pgxn/neon/libpagestore.c
@@ -581,6 +581,8 @@ pageserver_connect(shardno_t shard_no, int elevel)
 	}
 	case PS_Connecting_PageStream:
 	{
+		PGresult   *result;
+
 		neon_shard_log(shard_no, DEBUG5, "Connection state: Connecting_PageStream");
 
 		if (PQstatus(shard->conn) == CONNECTION_BAD)
@@ -620,6 +622,47 @@ pageserver_connect(shardno_t shard_no, int elevel)
 				}
 			}
 		}
+		result = PQgetResult(shard->conn);
+
+		if (!result)
+		{
+			CLEANUP_AND_DISCONNECT(shard);
+			neon_shard_log(shard_no, elevel,
+						   "could not complete handshake: unexpected end of data");
+			return false;
+		}
+
+		switch (PQresultStatus(result))
+		{
+		/*
+		 * We only expect COPY_BOTH, all other responses are indications of
+		 * problems out of our control
+		 */
+		case PGRES_COPY_BOTH:
+			break;
+		case PGRES_BAD_RESPONSE:
+		case PGRES_NONFATAL_ERROR:
+		case PGRES_FATAL_ERROR:
+			neon_shard_log(shard_no, elevel,
+						   "could not complete handshake: PageServer returned error: %s",
+						   PQresultErrorMessage(result));
+			PQclear(result);
+			return false;
+		case PGRES_EMPTY_QUERY:
+		case PGRES_COMMAND_OK:
+		case PGRES_TUPLES_OK:
+		case PGRES_COPY_OUT:
+		case PGRES_COPY_IN:
+		case PGRES_SINGLE_TUPLE:
+		case PGRES_PIPELINE_SYNC:
+		case PGRES_PIPELINE_ABORTED:
+			CLEANUP_AND_DISCONNECT(shard);
+			neon_shard_log(shard_no, elevel,
+						   "could not complete handshake: unexpected result type: %d",
+						   PQresultStatus(result));
+			PQclear(result);
+			return false;
+		}
 
 		shard->state = PS_Connected;
 		/* fallthrough */