Cancel task now has “last second” bail out check

impossibl · Oct 28, 2020 · 475691d · 475691d
1 parent b01be97
commit 475691d
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 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
@@ -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
@@ -53,6 +53,10 @@ enum State {
 
   protected abstract void go();
 
+  boolean isCancelled() {
+    return state.get() != State.Running || thread.isInterrupted();
+  }
+
   @Override
   public void run() {