From 0b4b9aa5a9b5e6aefc4f6d0dac872866b0157059 Mon Sep 17 00:00:00 2001
From: Jason Petersen <jason@citusdata.com>
Date: Wed, 19 Apr 2017 20:17:10 -0600
Subject: [PATCH] Add PG_WAIT_EXTENSION to Read/WriteFile calls

Read/WriteFile adds the same "wait_event_info" argument, so I've added
a PG_WAIT_EXTENSION to all these calls as well.
---
 src/backend/distributed/commands/transmit.c     | 17 +++++++++++++++++
 .../worker/worker_partition_protocol.c          |  6 ++++++
 2 files changed, 23 insertions(+)

diff --git a/src/backend/distributed/commands/transmit.c b/src/backend/distributed/commands/transmit.c
index af0fcf5e7..c9f01852e 100644
--- a/src/backend/distributed/commands/transmit.c
+++ b/src/backend/distributed/commands/transmit.c
@@ -9,6 +9,7 @@
 
 #include "postgres.h"
 #include "miscadmin.h"
+#include "pgstat.h"
 
 #include <fcntl.h>
 #include <sys/stat.h>
@@ -54,7 +55,13 @@ RedirectCopyDataToRegularFile(const char *filename)
 		/* if received data has contents, append to regular file */
 		if (copyData->len > 0)
 		{
+#if (PG_VERSION_NUM >= 100000)
+			int appended = FileWrite(fileDesc, copyData->data, copyData->len,
+									 PG_WAIT_EXTENSION);
+#else
 			int appended = FileWrite(fileDesc, copyData->data, copyData->len);
+#endif
+
 			if (appended != copyData->len)
 			{
 				ereport(ERROR, (errcode_for_file_access(),
@@ -98,7 +105,12 @@ SendRegularFile(const char *filename)
 
 	SendCopyOutStart();
 
+#if (PG_VERSION_NUM >= 100000)
+	readBytes = FileRead(fileDesc, fileBuffer->data, fileBufferSize, PG_WAIT_EXTENSION);
+#else
 	readBytes = FileRead(fileDesc, fileBuffer->data, fileBufferSize);
+#endif
+
 	while (readBytes > 0)
 	{
 		fileBuffer->len = readBytes;
@@ -106,7 +118,12 @@ SendRegularFile(const char *filename)
 		SendCopyData(fileBuffer);
 
 		resetStringInfo(fileBuffer);
+#if (PG_VERSION_NUM >= 100000)
+		readBytes = FileRead(fileDesc, fileBuffer->data, fileBufferSize,
+							 PG_WAIT_EXTENSION);
+#else
 		readBytes = FileRead(fileDesc, fileBuffer->data, fileBufferSize);
+#endif
 	}
 
 	SendCopyDone();
diff --git a/src/backend/distributed/worker/worker_partition_protocol.c b/src/backend/distributed/worker/worker_partition_protocol.c
index b15915f78..79ace362e 100644
--- a/src/backend/distributed/worker/worker_partition_protocol.c
+++ b/src/backend/distributed/worker/worker_partition_protocol.c
@@ -16,6 +16,7 @@
 
 #include "postgres.h"
 #include "funcapi.h"
+#include "pgstat.h"
 
 #include <arpa/inet.h>
 #include <netinet/in.h>
@@ -742,7 +743,12 @@ FileOutputStreamFlush(FileOutputStream file)
 	int written = 0;
 
 	errno = 0;
+#if (PG_VERSION_NUM >= 100000)
+	written = FileWrite(file.fileDescriptor, fileBuffer->data, fileBuffer->len,
+						PG_WAIT_EXTENSION);
+#else
 	written = FileWrite(file.fileDescriptor, fileBuffer->data, fileBuffer->len);
+#endif
 	if (written != fileBuffer->len)
 	{
 		ereport(ERROR, (errcode_for_file_access(),