ImmediateDieOk-1.patch

diff --git a/src/backend/libpq/be-secure.c b/src/backend/libpq/be-secure.c
index 41ec1ad..049e5b1 100644
--- a/src/backend/libpq/be-secure.c
+++ b/src/backend/libpq/be-secure.c
@@ -178,5 +178,13 @@ secure_write(Port *port, void *ptr, size_t len)
 ssize_t
 secure_raw_write(Port *port, const void *ptr, size_t len)
 {
-	return send(port->sock, ptr, len, 0);
+	ssize_t result;
+
+	prepare_for_client_write();
+
+	result = send(port->sock, ptr, len, 0);
+
+	client_write_ended();
+
+	return result;
 }
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
index 61f17bf..138060b 100644
--- a/src/backend/tcop/postgres.c
+++ b/src/backend/tcop/postgres.c
@@ -497,6 +497,17 @@ ReadCommand(StringInfo inBuf)
  * non-reentrant libc functions.  This restriction makes it safe for us
  * to allow interrupt service routines to execute nontrivial code while
  * we are waiting for input.
+ *
+ * When waiting in the main loop, we can process any interrupt immediately
+ * in the signal handler. In any other read from the client, like in a COPY
+ * FROM STDIN, we can't safely process a query cancel signal, because we might
+ * be in the middle of sending a message to the client, and jumping out would
+ * violate the protocol. Or rather, pqcomm.c would detect it and refuse to
+ * send any more messages to the client. But handling a SIGTERM is OK, because
+ * we're terminating the backend and don't need to send any more messages
+ * anyway. That means that we might not be able to send an error message to
+ * the client, but that seems better than waiting indefinitely, in case the
+ * client is not responding.
  */
 void
 prepare_for_client_read(void)
@@ -513,6 +524,15 @@ prepare_for_client_read(void)
 		/* And don't forget to detect one that already arrived */
 		CHECK_FOR_INTERRUPTS();
 	}
+	else
+	{
+		/* Allow die interrupts to be processed while waiting */
+		ImmediateDieOK = true;
+
+		/* And don't forget to detect one that already arrived */
+		if (ProcDiePending)
+			CHECK_FOR_INTERRUPTS();
+	}
 }
 
 /*
@@ -534,6 +554,35 @@ client_read_ended(void)
 
 		errno = save_errno;
 	}
+	else
+		ImmediateDieOK = false;
+}
+
+/*
+ * prepare_for_client_write -- set up to possibly block on client output
+ *
+ * Like prepare_client_read, but for writing.
+ */
+void
+prepare_for_client_write(void)
+{
+	/* Allow die interrupts to be processed while waiting */
+	ImmediateDieOK = true;
+
+	/* And don't forget to detect one that already arrived */
+	if (ProcDiePending)
+		CHECK_FOR_INTERRUPTS();
+}
+
+/*
+ * client_read_ended -- get out of the client-output state
+ *
+ * This is called just after low-level writes.
+ */
+void
+client_write_ended(void)
+{
+	ImmediateDieOK = false;
 }
 
 
@@ -2591,8 +2640,8 @@ die(SIGNAL_ARGS)
 		 * If it's safe to interrupt, and we're waiting for input or a lock,
 		 * service the interrupt immediately
 		 */
-		if (ImmediateInterruptOK && InterruptHoldoffCount == 0 &&
-			CritSectionCount == 0)
+		if ((ImmediateInterruptOK || ImmediateDieOK) &&
+			InterruptHoldoffCount == 0 && CritSectionCount == 0)
 		{
 			/* bump holdoff count to make ProcessInterrupts() a no-op */
 			/* until we are done getting ready for it */
@@ -2792,8 +2841,8 @@ RecoveryConflictInterrupt(ProcSignalReason reason)
 		 * If it's safe to interrupt, and we're waiting for input or a lock,
 		 * service the interrupt immediately
 		 */
-		if (ImmediateInterruptOK && InterruptHoldoffCount == 0 &&
-			CritSectionCount == 0)
+		if ((ImmediateInterruptOK || (ImmediateDieOK && ProcDiePending)) &&
+			InterruptHoldoffCount == 0 && CritSectionCount == 0)
 		{
 			/* bump holdoff count to make ProcessInterrupts() a no-op */
 			/* until we are done getting ready for it */
diff --git a/src/backend/utils/init/globals.c b/src/backend/utils/init/globals.c
index be74835..24523f9 100644
--- a/src/backend/utils/init/globals.c
+++ b/src/backend/utils/init/globals.c
@@ -31,6 +31,7 @@ volatile bool QueryCancelPending = false;
 volatile bool ProcDiePending = false;
 volatile bool ClientConnectionLost = false;
 volatile bool ImmediateInterruptOK = false;
+volatile bool ImmediateDieOK = false;
 volatile uint32 InterruptHoldoffCount = 0;
 volatile uint32 CritSectionCount = 0;
 
diff --git a/src/include/miscadmin.h b/src/include/miscadmin.h
index 2ba9885..9528936 100644
--- a/src/include/miscadmin.h
+++ b/src/include/miscadmin.h
@@ -81,6 +81,7 @@ extern volatile bool ClientConnectionLost;
 
 /* these are marked volatile because they are examined by signal handlers: */
 extern PGDLLIMPORT volatile bool ImmediateInterruptOK;
+extern PGDLLIMPORT volatile bool ImmediateDieOK;
 extern PGDLLIMPORT volatile uint32 InterruptHoldoffCount;
 extern PGDLLIMPORT volatile uint32 CritSectionCount;
 
diff --git a/src/include/tcop/tcopprot.h b/src/include/tcop/tcopprot.h
index 60f7532..c288bdd 100644
--- a/src/include/tcop/tcopprot.h
+++ b/src/include/tcop/tcopprot.h
@@ -69,6 +69,8 @@ extern void RecoveryConflictInterrupt(ProcSignalReason reason); /* called from S
 																 * handler */
 extern void prepare_for_client_read(void);
 extern void client_read_ended(void);
+extern void prepare_for_client_write(void);
+extern void client_write_ended(void);
 extern void process_postgres_switches(int argc, char *argv[],
 						  GucContext ctx, const char **dbname);
 extern void PostgresMain(int argc, char *argv[],