Batch rows on access node for distributed COPY

src/guc.c

@@ -62,6 +62,13 @@ static const struct config_enum_entry hypertable_distributed_types[] = {
 	{ NULL, 0, false }
 };
 
+static const struct config_enum_entry dist_copy_transfer_formats[] = {
+	{ "auto", DCTF_Auto, false },
+	{ "binary", DCTF_Binary, false },
+	{ "text", DCTF_Text, false },
+	{ NULL, 0, false }
+};
+
 bool ts_guc_enable_optimizations = true;
 bool ts_guc_restoring = false;
 bool ts_guc_enable_constraint_aware_append = true;
@@ -90,6 +97,7 @@ char *ts_last_tune_version = NULL;
 TSDLLEXPORT bool ts_guc_enable_2pc;
 TSDLLEXPORT int ts_guc_max_insert_batch_size = 1000;
 TSDLLEXPORT bool ts_guc_enable_connection_binary_data;
+TSDLLEXPORT DistCopyTransferFormat ts_guc_dist_copy_transfer_format;
 TSDLLEXPORT bool ts_guc_enable_client_ddl_on_data_nodes = false;
 TSDLLEXPORT char *ts_guc_ssl_dir = NULL;
 TSDLLEXPORT char *ts_guc_passfile = NULL;
@@ -309,6 +317,25 @@ _guc_init(void)
 							 NULL,
 							 NULL);
 
+	/*
+	 * The default is 'auto', so that the dist COPY could use text transfer
+	 * format for text input. It has a passthrough optimization for this case,
+	 * which greatly reduces the CPU usage. Ideally we would implement the same
+	 * optimization for binary, but the Postgres COPY code doesn't provide
+	 * enough APIs for that.
+	 */
+	DefineCustomEnumVariable("timescaledb.dist_copy_transfer_format",
+							 "Data format used by distributed COPY to send data to data nodes",
+							 "auto, binary or text",
+							 (int *) &ts_guc_dist_copy_transfer_format,
+							 DCTF_Auto,
+							 dist_copy_transfer_formats,
+							 PGC_USERSET,
+							 0,
+							 NULL,
+							 NULL,
+							 NULL);
+
 	DefineCustomBoolVariable("timescaledb.enable_client_ddl_on_data_nodes",
 							 "Enable DDL operations on data nodes by a client",
 							 "Do not restrict execution of DDL operations only by access node",

src/guc.h

@@ -66,6 +66,15 @@ typedef enum HypertableDistType
 extern TSDLLEXPORT HypertableDistType ts_guc_hypertable_distributed_default;
 extern TSDLLEXPORT int ts_guc_hypertable_replication_factor_default;
 
+typedef enum DistCopyTransferFormat
+{
+	DCTF_Auto,
+	DCTF_Binary,
+	DCTF_Text
+} DistCopyTransferFormat;
+
+extern TSDLLEXPORT DistCopyTransferFormat ts_guc_dist_copy_transfer_format;
+
 #ifdef TS_DEBUG
 extern bool ts_shutdown_bgw;
 extern char *ts_current_timestamp_mock;

test/src/test_utils.c

@@ -16,6 +16,7 @@
 #include <storage/procarray.h>
 #include <utils/builtins.h>
 #include <utils/elog.h>
+#include <utils/guc.h>
 #include <utils/memutils.h>
 
 #include "debug_point.h"
@@ -125,19 +126,34 @@ ts_debug_shippable_fatal_after_n_rows(PG_FUNCTION_ARGS)
 
 /*
  * Broken send/receive functions for int4 that throw after an (arbitrarily
- * chosen prime) number of rows.
- * Use ERROR, not FATAL, because PG versions < 14 are unable to report a FATAL
- * error to the access node before closing the connection, so the test results
- * would be different.
+ * chosen prime or configured) number of rows.
  */
-#define ARBITRARY_PRIME_NUMBER 7103
+static void
+broken_sendrecv_throw()
+{
+	int throw_after = 7103; /* an arbitrary prime */
+	const char *throw_after_option =
+		GetConfigOption("timescaledb.debug_broken_sendrecv_throw_after", true, false);
+
+	if (throw_after_option)
+	{
+		throw_after = pg_atoi(throw_after_option, 4, 0);
+	}
+
+	/*
+	 * Use ERROR, not FATAL, because PG versions < 14 are unable to report a
+	 * FATAL error to the access node before closing the connection, so the test
+	 * results would be different.
+	 */
+	(void) throw_after_n_rows(throw_after, ERROR);
+}
 
 TS_FUNCTION_INFO_V1(ts_debug_broken_int4recv);
 
 Datum
 ts_debug_broken_int4recv(PG_FUNCTION_ARGS)
 {
-	(void) throw_after_n_rows(ARBITRARY_PRIME_NUMBER, ERROR);
+	broken_sendrecv_throw();
 	return int4recv(fcinfo);
 }
 
@@ -146,7 +162,7 @@ TS_FUNCTION_INFO_V1(ts_debug_broken_int4send);
 Datum
 ts_debug_broken_int4send(PG_FUNCTION_ARGS)
 {
-	(void) throw_after_n_rows(ARBITRARY_PRIME_NUMBER, ERROR);
+	broken_sendrecv_throw();
 	return int4send(fcinfo);
 }
 

tsl/src/nodes/data_node_copy.c

@@ -288,7 +288,7 @@ data_node_copy_end(CustomScanState *node)
 	DataNodeCopyState *dncs = (DataNodeCopyState *) node;
 
 	ExecEndNode(linitial(node->custom_ps));
-	remote_copy_end(dncs->copy_ctx);
+	remote_copy_end_on_success(dncs->copy_ctx);
 	ts_cache_release(dncs->hcache);
 }
 

tsl/src/remote/connection.c

@@ -847,6 +847,9 @@ remote_connection_set_status(TSConnection *conn, TSConnectionStatus status)
 {
 	Assert(conn != NULL);
 	conn->status = status;
+
+	/* Should be blocking except when doing COPY. */
+	Assert(PQisnonblocking(conn->pg_conn) == (conn->status == CONN_COPY_IN));
 }
 
 TSConnectionStatus
@@ -2162,6 +2165,26 @@ remote_connection_begin_copy(TSConnection *conn, const char *copycmd, bool binar
 								 "connection not IDLE when beginning COPY",
 								 conn);
 
+#ifndef NDEBUG
+	/* Set some variables for testing. */
+	const char *throw_after_option =
+		GetConfigOption("timescaledb.debug_broken_sendrecv_throw_after", true, false);
+	if (throw_after_option)
+	{
+		res = PQexec(pg_conn,
+					 psprintf("set timescaledb.debug_broken_sendrecv_throw_after = '%s';",
+							  throw_after_option));
+		if (PQresultStatus(res) != PGRES_COMMAND_OK)
+		{
+			remote_connection_get_result_error(res, err);
+			PQclear(res);
+			return false;
+		}
+		PQclear(res);
+	}
+#endif
+
+	/* Run the COPY query. */
 	res = PQexec(pg_conn, copycmd);
 
 	if (PQresultStatus(res) != PGRES_COPY_IN)
@@ -2237,8 +2260,7 @@ send_end_binary_copy_data(const TSConnection *conn, TSConnectionError *err)
 bool
 remote_connection_end_copy(TSConnection *conn, TSConnectionError *err)
 {
-	PGresult *res;
-	bool success;
+	PGresult *res = NULL;
 
 	/*
 	 * In any case, try to switch the connection into the blocking mode, because
@@ -2331,21 +2353,26 @@ remote_connection_end_copy(TSConnection *conn, TSConnectionError *err)
 	if (res == NULL || PQresultStatus(res) != PGRES_COPY_IN)
 	{
 		remote_connection_set_status(conn, res == NULL ? CONN_IDLE : CONN_PROCESSING);
-		elog(ERROR, "connection marked as CONN_COPY_IN, but no COPY is in progress");
 	}
 
-	if (conn->binary_copy && !send_end_binary_copy_data(conn, err))
-		return false;
+	/*
+	 * Finish the COPY if needed.
+	 */
+	if (remote_connection_get_status(conn) == CONN_COPY_IN)
+	{
+		if (conn->binary_copy && !send_end_binary_copy_data(conn, err))
+			return false;
 
-	if (PQputCopyEnd(conn->pg_conn, NULL) != 1)
-		return fill_connection_error(err,
-									 ERRCODE_CONNECTION_EXCEPTION,
-									 "could not end remote COPY",
-									 conn);
+		if (PQputCopyEnd(conn->pg_conn, NULL) != 1)
+			return fill_connection_error(err,
+										 ERRCODE_CONNECTION_EXCEPTION,
+										 "could not end remote COPY",
+										 conn);
 
-	success = true;
-	remote_connection_set_status(conn, CONN_PROCESSING);
+		remote_connection_set_status(conn, CONN_PROCESSING);
+	}
 
+	bool success = true;
 	while ((res = PQgetResult(conn->pg_conn)))
 	{
 		ExecStatusType status = PQresultStatus(res);