From f08de1c11de2cca22d2caadc368d3ef1e000e39a Mon Sep 17 00:00:00 2001
From: Matt Corallo <git@bluematt.me>
Date: Mon, 1 Feb 2021 20:26:00 -0500
Subject: [PATCH] Small tweaks to NioPeerHandler

---
 src/main/java/org/ldk/batteries/NioPeerHandler.java | 7 ++++++-
 src/test/java/org/ldk/HumanObjectPeerTest.java      | 6 +++++-
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/src/main/java/org/ldk/batteries/NioPeerHandler.java b/src/main/java/org/ldk/batteries/NioPeerHandler.java
index 29949622..14c5611d 100644
--- 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(() -> {
diff --git a/src/test/java/org/ldk/HumanObjectPeerTest.java b/src/test/java/org/ldk/HumanObjectPeerTest.java
index 27ddc7c2..807b3307 100644
--- a/src/test/java/org/ldk/HumanObjectPeerTest.java
+++ b/src/test/java/org/ldk/HumanObjectPeerTest.java
@@ -379,7 +379,7 @@ class HumanObjectPeerTestInstance {
         if (use_nio_peer_handler) {
             peer1.nio_peer_handler.check_events();
             peer2.nio_peer_handler.check_events();
-            try { Thread.sleep(500); } catch (InterruptedException e) { assert false; }
+            try { Thread.sleep(400); } catch (InterruptedException e) { assert false; }
         } else {
             synchronized (runqueue) {
                 ran = false;
@@ -610,6 +610,10 @@ class HumanObjectPeerTestInstance {
         }
 
         if (use_nio_peer_handler) {
+            state.peer1.peer_manager.disconnect_by_node_id(state.peer2.chan_manager.get_our_node_id(), false);
+            wait_events_processed(state.peer1, state.peer2);
+            assert state.peer1.peer_manager.get_peer_node_ids().length == 0;
+            assert state.peer2.peer_manager.get_peer_node_ids().length == 0;
             state.peer1.nio_peer_handler.interrupt();
             state.peer2.nio_peer_handler.interrupt();
         }
-- 
2.30.2