Add a timeout to NioPeerHandler.connect

[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 299496229ad73980914ad71ce282fdde1e0842b5..a77f1f0f49bb8cb4152d284b09f05c6b0f068116 100644 (file)
--- a/src/main/java/org/ldk/batteries/NioPeerHandler.java
+++ b/src/main/java/org/ldk/batteries/NioPeerHandler.java
@@ -11,6 +11,10 @@ import java.nio.channels.Selector;
 import java.nio.channels.ServerSocketChannel;
 import java.nio.channels.SocketChannel;
 
+/**
+ * A NioPeerHandler maps LDK's PeerHandler to Java's NIO I/O interface. It spawns a single background thread which
+ * processes socket events and provides the data to LDK for decryption and processing.
+ */
 public class NioPeerHandler {
     private static class Peer {
         SocketDescriptor descriptor;
@@ -39,11 +43,13 @@ public class NioPeerHandler {
             public long send_data(byte[] data, boolean resume_read) {
                 if (resume_read) {
                     peer.key.interestOps(peer.key.interestOps() | SelectionKey.OP_READ);
+                    selector.wakeup();
                 }
                 try {
                     long written = chan.write(ByteBuffer.wrap(data));
                     if (written != data.length) {
                         peer.key.interestOps(peer.key.interestOps() | SelectionKey.OP_WRITE);
+                        selector.wakeup();
                     }
                     return written;
                 } catch (IOException e) {
@@ -88,7 +94,6 @@ public class NioPeerHandler {
      * @throws IOException If an internal java.nio error occurs.
      */
     public NioPeerHandler(PeerManager manager) throws IOException {
-long id = manager._test_only_get_ptr();
         this.peer_manager = manager;
         this.selector = Selector.open();
         io_thread = new Thread(() -> {
@@ -184,10 +189,20 @@ long id = manager._test_only_get_ptr();
      * @param their_node_id A valid 33-byte public key representing the peer's Lightning Node ID. If this is invalid,
      *                      undefined behavior (read: Segfault, etc) may occur.
      * @param remote The socket address to connect to.
+     * @param timeout_ms The amount of time, in milliseconds, up to which we will wait for connection to complete.
      * @throws IOException If connecting to the remote endpoint fails or internal java.nio errors occur.
      */
-    public void connect(byte[] their_node_id, SocketAddress remote) throws IOException {
-        SocketChannel chan = SocketChannel.open(remote);
+    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);
+        }
+        if (!chan.finishConnect()) { // Note that this may throw its own IOException if we failed for another reason
+            throw new IOException("Timed out");
+        }
         Peer peer = setup_socket(chan);
         Result_CVec_u8ZPeerHandleErrorZ res = this.peer_manager.new_outbound_connection(their_node_id, peer.descriptor);
         if (res instanceof  Result_CVec_u8ZPeerHandleErrorZ.Result_CVec_u8ZPeerHandleErrorZ_OK) {