track connection state

Author: serprex

src/binary.c

@@ -316,6 +316,14 @@ struct ch_binary_state
 
 	/* Per-query cancel callback (no userdata). */
 	bool		(*check_cancel_fn) (void);
+
+	/*
+	 * Set when an unrecoverable protocol/IO error happened (eg server
+	 * raised an exception mid-INSERT and closed the socket, or a write
+	 * hit EPIPE). Callers must drop the cached connection rather than
+	 * reuse it; the cache layer checks this via ch_binary_is_broken.
+	 */
+	bool		broken;
 };
 
 /* ch_binary_connection_t is defined in internal.h with three fields; we
@@ -515,6 +523,16 @@ ch_binary_connect(ch_connection_details * details, char *errbuf, size_t errbuf_s
 	return NULL;
 }
 
+bool
+ch_binary_is_broken(const ch_binary_connection_t * conn)
+{
+	if (!conn)
+		return false;
+	const struct ch_binary_state *s = (const struct ch_binary_state *) conn->client;
+
+	return s ? s->broken : false;
+}
+
 void
 ch_binary_close(ch_binary_connection_t * conn)
 {
@@ -808,6 +826,7 @@ ch_binary_simple_query(ch_binary_connection_t * conn, const ch_query * query,
 	if (rc != CHC_OK)
 	{
 		resp_set_error(resp, err.msg[0] ? err.msg : "send_query failed");
+		s->broken = true;
 		goto done;
 	}
 
@@ -823,6 +842,7 @@ ch_binary_simple_query(ch_binary_connection_t * conn, const ch_query * query,
 		if (rc != CHC_OK)
 		{
 			resp_set_error(resp, err.msg[0] ? err.msg : "recv_packet failed");
+			s->broken = true;
 			break;
 		}
 		if (check_cancel && check_cancel() && !resp->error)
@@ -854,6 +874,14 @@ ch_binary_simple_query(ch_binary_connection_t * conn, const ch_query * query,
 			case CHC_PKT_EXCEPTION:
 				resp_set_exception(resp, pkt.exception);
 				chc_packet_clear(s->client, &pkt);
+				/*
+				 * Older servers (and some protocol states) close the socket
+				 * after raising an exception, so reusing this connection for
+				 * a follow-up query risks EPIPE. Match the legacy C++ driver
+				 * (which always called Client::ResetConnection) and treat
+				 * the connection as broken.
+				 */
+				s->broken = true;
 				goto done;
 			case CHC_PKT_END_OF_STREAM:
 				chc_packet_clear(s->client, &pkt);
@@ -1025,6 +1053,8 @@ struct ch_binary_insert_handle
 {
 	chc_client *client;
 	chc_alloc	alloc;
+	struct ch_binary_state *state;	/* parent connection; used to flag
+									 * broken state on error */
 	chc_block  *initial_block;	/* schema source (server's empty Data) */
 	chc_query_opts opts;
 	size_t		ncols;
@@ -1197,6 +1227,7 @@ recv_initial_block(struct ch_binary_state *s, ch_binary_insert_handle * h,
 		{
 			set_errbuf(errbuf, errbuf_size,
 					   err.msg[0] ? err.msg : "recv_packet failed");
+			s->broken = true;
 			return -1;
 		}
 		if (pkt.kind == CHC_PKT_EXCEPTION)
@@ -1208,6 +1239,7 @@ recv_initial_block(struct ch_binary_state *s, ch_binary_insert_handle * h,
 			else if (pkt.exception && pkt.exception->name)
 				msg = pkt.exception->name;
 			set_errbuf(errbuf, errbuf_size, msg);
+			s->broken = true;
 			chc_packet_clear(s->client, &pkt);
 			return -1;
 		}
@@ -1238,6 +1270,7 @@ ch_binary_begin_insert(ch_binary_connection_t * conn, const ch_query * query,
 	}
 	h->client = s->client;
 	h->alloc = s->alloc;
+	h->state = s;
 	arena_init(&h->info_arena);
 
 	/* Append " VALUES" so the server enters insert mode. */
@@ -1284,6 +1317,7 @@ ch_binary_begin_insert(ch_binary_connection_t * conn, const ch_query * query,
 	if (rc != CHC_OK)
 	{
 		set_errbuf(errbuf, errbuf_size, err.msg[0] ? err.msg : "send_query failed");
+		s->broken = true;
 		ch_binary_end_insert(h, errbuf, errbuf_size);
 		return NULL;
 	}
@@ -2219,6 +2253,8 @@ ch_binary_flush_block(ch_binary_insert_handle * h, char *errbuf, size_t errbuf_s
 	if (rc != CHC_OK)
 	{
 		set_errbuf(errbuf, errbuf_size, err.msg[0] ? err.msg : "send_data");
+		if (h->state)
+			h->state->broken = true;
 		goto cleanup;
 	}
 	ret = 0;
@@ -2271,6 +2307,8 @@ ch_binary_end_insert(ch_binary_insert_handle * h, char *errbuf, size_t errbuf_si
 				{
 					set_errbuf(errbuf, errbuf_size,
 							   err.msg[0] ? err.msg : "recv_packet failed");
+					if (h->state)
+						h->state->broken = true;
 					chc_packet_clear(h->client, &pkt);
 					break;
 				}
@@ -2283,6 +2321,13 @@ ch_binary_end_insert(ch_binary_insert_handle * h, char *errbuf, size_t errbuf_si
 					else if (pkt.exception && pkt.exception->name)
 						msg = pkt.exception->name;
 					set_errbuf(errbuf, errbuf_size, msg);
+					/*
+					 * Server raised an exception mid-INSERT and typically
+					 * closes the socket; the next op on this conn would
+					 * EPIPE. Mark broken so the cache rebuilds.
+					 */
+					if (h->state)
+						h->state->broken = true;
 					chc_packet_clear(h->client, &pkt);
 					break;
 				}
@@ -2294,6 +2339,8 @@ ch_binary_end_insert(ch_binary_insert_handle * h, char *errbuf, size_t errbuf_si
 		else
 		{
 			set_errbuf(errbuf, errbuf_size, err.msg[0] ? err.msg : "send_data terminator failed");
+			if (h->state)
+				h->state->broken = true;
 		}
 	}
 	if (h->initial_block)

src/connection.c

@@ -123,10 +123,19 @@ chfdw_get_connection(UserMapping * user)
 	}
 
 	/*
-	 * We don't check the health of cached connection here, because it would
-	 * require some overhead. Broken connection will be detected when the
-	 * connection is actually used.
+	 * Drop connections that hit an unrecoverable protocol/IO error on the
+	 * previous statement (server raised mid-INSERT and closed the socket,
+	 * write hit EPIPE, etc). Without this, a subsequent statement would
+	 * write to the dead socket and surface a useless "Broken pipe" instead
+	 * of the real error from the next request.
 	 */
+	if (entry->gate.conn != NULL && entry->gate.methods->is_broken != NULL
+		&& entry->gate.methods->is_broken(entry->gate.conn))
+	{
+		elog(LOG, "closing broken pg_clickhouse connection");
+		entry->gate.methods->disconnect(entry->gate.conn);
+		entry->gate.conn = NULL;
+	}
 
 	/*
 	 * If cache entry doesn't have a connection, we have to establish a new

src/include/binary.h

@@ -41,6 +41,13 @@ extern ch_binary_connection_t * ch_binary_connect(ch_connection_details * detail
 												  char *errbuf, size_t errbuf_size);
 extern void ch_binary_close(ch_binary_connection_t * conn);
 
+/*
+ * Returns true if the connection encountered an unrecoverable error
+ * (server exception, IO failure, mid-protocol break). Callers should
+ * drop the cached connection instead of reusing it.
+ */
+extern bool ch_binary_is_broken(const ch_binary_connection_t * conn);
+
 /* SELECT. */
 extern ch_binary_response_t * ch_binary_simple_query(ch_binary_connection_t * conn,
 													 const ch_query * query,

src/include/fdw.h

@@ -78,6 +78,7 @@ typedef void (*insert_tuple_method) (void *state, TupleTableSlot * slot);
 typedef ch_cursor * (*streaming_query_method) (void *conn,
 											   const ch_query * query,
 											   int32 fetch_size);
+typedef bool (*is_broken_method) (const void *conn);
 
 typedef struct
 {
@@ -88,6 +89,7 @@ typedef struct
 	insert_tuple_method insert_tuple;
 	streaming_query_method streaming_query; /* NULL if not supported */
 	cursor_fetch_row_method streaming_fetch_row;	/* NULL if not supported */
+	is_broken_method is_broken; /* NULL means connection is never broken */
 }			libclickhouse_methods;
 
 typedef struct

src/pglink.c

@@ -58,6 +58,7 @@ static libclickhouse_methods http_methods =
 static void binary_disconnect(void *conn);
 static ch_cursor * binary_simple_query(void *conn, const ch_query * query);
 static void binary_cursor_free(void *cursor);
+static bool binary_is_broken(const void *conn);
 
 /* static void binary_simple_insert(void *conn, const char *query); */
 static Datum * binary_fetch_row(ChFdwScanRowContext * ctx);
@@ -78,6 +79,7 @@ static libclickhouse_methods binary_methods =
 		.insert_tuple = binary_insert_tuple,
 		.streaming_query = NULL,
 		.streaming_fetch_row = NULL,
+		.is_broken = binary_is_broken,
 };
 
 static int
@@ -841,6 +843,12 @@ binary_disconnect(void *conn)
 		ch_binary_close((ch_binary_connection_t *) conn);
 }
 
+static bool
+binary_is_broken(const void *conn)
+{
+	return ch_binary_is_broken((const ch_binary_connection_t *) conn);
+}
+
 static ch_cursor *
 binary_simple_query(void *conn, const ch_query * query)
 {