From 475691d6f4037d3821e5cf5ad291daa0214414b7 Mon Sep 17 00:00:00 2001
From: Kevin Wooten <kevin@wooten.com>
Date: Wed, 28 Oct 2020 11:52:00 -0700
Subject: [PATCH] =?UTF-8?q?Cancel=20task=20now=20has=20=E2=80=9Clast=20sec?=
 =?UTF-8?q?ond=E2=80=9D=20bail=20out=20check?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../postgres/jdbc/CancelRequestTask.java      | 26 ++++++++++++++++++-
 .../postgres/jdbc/ExecutionTimerTask.java     |  4 +++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/driver/src/main/java/com/impossibl/postgres/jdbc/CancelRequestTask.java b/driver/src/main/java/com/impossibl/postgres/jdbc/CancelRequestTask.java
index afbbbcc59..066b5b397 100644
--- a/driver/src/main/java/com/impossibl/postgres/jdbc/CancelRequestTask.java
+++ b/driver/src/main/java/com/impossibl/postgres/jdbc/CancelRequestTask.java
@@ -40,6 +40,7 @@
 import java.io.ByteArrayOutputStream;
 import java.io.DataOutputStream;
 import java.io.IOException;
+import java.io.OutputStream;
 import java.net.InetSocketAddress;
 import java.net.Socket;
 import java.net.SocketAddress;
@@ -67,6 +68,10 @@ public void go() {
 
   private void sendCancelRequest() {
 
+    if (isCancelled()) {
+      return;
+    }
+
     if (keyData.getProcessId() == 0 && keyData.getSecretKey() == 0) {
       logger.warning("Cannot send CancelRequest because of missing BackendKeyData.");
       return;
@@ -79,7 +84,21 @@ private void sendCancelRequest() {
         InetSocketAddress target = (InetSocketAddress) serverAddress;
 
         try (Socket abortSocket = new Socket(target.getAddress(), target.getPort())) {
-          writeCancelRequest(new DataOutputStream(abortSocket.getOutputStream()));
+          OutputStream abortSocketStream = abortSocket.getOutputStream();
+
+          try (ByteArrayOutputStream os = new ByteArrayOutputStream()) {
+
+            writeCancelRequest(new DataOutputStream(os));
+
+            byte[] request = os.toByteArray();
+
+            // last second bail out
+            if (isCancelled()) {
+              return;
+            }
+
+            abortSocketStream.write(request);
+          }
         }
 
       }
@@ -98,6 +117,11 @@ else if (serverAddress instanceof DomainSocketAddress) {
             buffer.put(request, 0, request.length);
             buffer.flip();
 
+            // last second bail out
+            if (isCancelled()) {
+              return;
+            }
+
             unixSocket.write(buffer, 0, buffer.limit());
           }
 
diff --git a/driver/src/main/java/com/impossibl/postgres/jdbc/ExecutionTimerTask.java b/driver/src/main/java/com/impossibl/postgres/jdbc/ExecutionTimerTask.java
index 883e3059a..1a1f39d16 100644
--- a/driver/src/main/java/com/impossibl/postgres/jdbc/ExecutionTimerTask.java
+++ b/driver/src/main/java/com/impossibl/postgres/jdbc/ExecutionTimerTask.java
@@ -53,6 +53,10 @@ enum State {
 
   protected abstract void go();
 
+  boolean isCancelled() {
+    return state.get() != State.Running || thread.isInterrupted();
+  }
+
   @Override
   public void run() {