[Java] More carefully ensure sockets are closed in `NioPeerHandler`

[ldk-java] / src / main / java / org / ldk / batteries / NioPeerHandler.java

diff --git a/src/main/java/org/ldk/batteries/NioPeerHandler.java b/src/main/java/org/ldk/batteries/NioPeerHandler.java
index 49f60c7f721b22f570dd5e5d9c300ca9c9503fbe..2ab48d627a67c0405d985cd6fdaf78969b276ff1 100644 (file)
--- a/src/main/java/org/ldk/batteries/NioPeerHandler.java
+++ b/src/main/java/org/ldk/batteries/NioPeerHandler.java
@@ -5,6 +5,7 @@ import org.ldk.structs.*;
 
 import java.io.IOException;
 import java.lang.reflect.Field;
+import java.lang.ref.Reference;
 import java.util.LinkedList;
 import java.net.SocketAddress;
 import java.net.StandardSocketOptions;
@@ -86,7 +87,7 @@ public class NioPeerHandler {
                                                        (peer.key.interestOps() | SelectionKey.OP_READ) & (~SelectionKey.OP_WRITE)));
                                        }
                     return written;
-                } catch (IOException e) {
+                } catch (IOException|CancelledKeyException ignored) {
                     // Most likely the socket is disconnected, let the background thread handle it.
                     return 0;
                 }
@@ -96,7 +97,7 @@ public class NioPeerHandler {
             public void disconnect_socket() {
                 try {
                     do_selector_action(() -> {
-                        peer.key.cancel();
+                        try { peer.key.cancel(); } catch (CancelledKeyException ignored) {}
                         peer.key.channel().close();
                     });
                 } catch (IOException ignored) { }
@@ -188,8 +189,8 @@ public class NioPeerHandler {
                             if (key.isValid() && (key.interestOps() & SelectionKey.OP_WRITE) != 0 && key.isWritable()) {
                                 Result_NonePeerHandleErrorZ res = this.peer_manager.write_buffer_space_avail(peer.descriptor);
                                 if (res instanceof Result_NonePeerHandleErrorZ.Result_NonePeerHandleErrorZ_Err) {
-                                    key.channel().close();
                                     key.cancel();
+                                    key.channel().close();
                                 }
                             }
                             if (key.isValid() && (key.interestOps() & SelectionKey.OP_READ) != 0 && key.isReadable()) {
@@ -198,6 +199,7 @@ public class NioPeerHandler {
                                 if (read == -1) {
                                     this.peer_manager.socket_disconnected(peer.descriptor);
                                     key.cancel();
+                                    key.channel().close(); // This may throw, we read -1 so the channel should already be closed, but do this to be safe
                                 } else if (read > 0) {
                                     ((Buffer)buf).flip();
                                     // This code is quite hot during initial network graph sync, so we go a ways out of
@@ -218,19 +220,19 @@ public class NioPeerHandler {
                                     long read_result_pointer = bindings.PeerManager_read_event(
                                             peer_manager_raw_pointer, peer.descriptor_raw_pointer, read_bytes);
                                     if (bindings.CResult_boolPeerHandleErrorZ_is_ok(read_result_pointer)) {
-                                        if (bindings.LDKCResult_boolPeerHandleErrorZ_get_ok(read_result_pointer)) {
+                                        if (bindings.CResult_boolPeerHandleErrorZ_get_ok(read_result_pointer)) {
                                             key.interestOps(key.interestOps() & (~SelectionKey.OP_READ));
                                         }
                                     } else {
-                                        key.channel().close();
                                         key.cancel();
+                                        key.channel().close();
                                     }
                                     bindings.CResult_boolPeerHandleErrorZ_free(read_result_pointer);
                                 }
                             }
                         } catch (IOException ignored) {
-                            try { key.channel().close(); } catch (IOException ignored2) { }
                             key.cancel();
+                            try { key.channel().close(); } catch (IOException ignored2) { }
                             peer_manager.socket_disconnected(peer.descriptor);
                         }
                     } catch (CancelledKeyException e) {
@@ -264,13 +266,21 @@ public class NioPeerHandler {
      */
     public void connect(byte[] their_node_id, SocketAddress remote, int timeout_ms) throws IOException {
         SocketChannel chan = SocketChannel.open();
-        chan.configureBlocking(false);
-        Selector open_selector = Selector.open();
-        chan.register(open_selector, SelectionKey.OP_CONNECT);
-        if (!chan.connect(remote)) {
-            open_selector.select(timeout_ms);
+        boolean connected;
+        try {
+            chan.configureBlocking(false);
+            Selector open_selector = Selector.open();
+            chan.register(open_selector, SelectionKey.OP_CONNECT);
+            if (!chan.connect(remote)) {
+                open_selector.select(timeout_ms);
+            }
+            connected = chan.finishConnect();
+        } catch (IOException e) {
+            try { chan.close(); } catch (IOException _e) { }
+            throw e;
         }
-        if (!chan.finishConnect()) { // Note that this may throw its own IOException if we failed for another reason
+        if (!connected) {
+            try { chan.close(); } catch (IOException _e) { }
             throw new IOException("Timed out");
         }
         Peer peer = setup_socket(chan);
@@ -339,6 +349,7 @@ public class NioPeerHandler {
                 }
             } catch (IOException ignored) {}
         }
+        Reference.reachabilityFence(this.peer_manager); // Almost certainly overkill, but no harm in it
     }
 
     /**