Fix premature close connection

Artem Labazin · Artem Labazin · commit 196dc2fe9c6a · 2018-06-26T13:57:48.000+03:00
diff --git a/encon/src/main/java/io/appulse/encon/ModuleClient.java b/encon/src/main/java/io/appulse/encon/ModuleClient.java
@@ -28,7 +28,6 @@
 import io.appulse.encon.common.RemoteNode;
 import io.appulse.encon.connection.Connection;
 import io.appulse.encon.connection.handshake.HandshakeClientInitializer;
-import io.appulse.encon.connection.regular.ConnectionHandler;
 
 import io.netty.bootstrap.Bootstrap;
 import lombok.NonNull;
@@ -94,17 +93,7 @@ private CompletableFuture<Connection> createConnection (@NonNull RemoteNode remo
             .node(node)
             .future(future)
             .remote(remote)
-            .channelCloseListener(f -> {
-              log.debug("Running close listener");
-              ConnectionHandler connectionHandler = f.channel()
-                  .pipeline()
-                  .get(ConnectionHandler.class);
-
-              if (connectionHandler == null) {
-                return;
-              }
-
-              RemoteNode remoteNode = connectionHandler.getRemote();
+            .channelCloseAction(remoteNode -> {
               node.moduleLookup.remove(remoteNode);
               node.moduleConnection.remove(remoteNode);
             })
diff --git a/encon/src/main/java/io/appulse/encon/ModuleServer.java b/encon/src/main/java/io/appulse/encon/ModuleServer.java
@@ -27,9 +27,7 @@
 
 import java.io.Closeable;
 
-import io.appulse.encon.common.RemoteNode;
 import io.appulse.encon.connection.handshake.HandshakeServerInitializer;
-import io.appulse.encon.connection.regular.ConnectionHandler;
 
 import io.netty.bootstrap.ServerBootstrap;
 import io.netty.buffer.PooledByteBufAllocator;
@@ -81,17 +79,7 @@ private void start () {
         .childHandler(HandshakeServerInitializer.builder()
             .node(node)
             .consumer(moduleConnection::add)
-            .channelCloseListener(future -> {
-              log.debug("Running close listener");
-              ConnectionHandler connectionHandler = future.channel()
-                  .pipeline()
-                  .get(ConnectionHandler.class);
-
-              if (connectionHandler == null) {
-                return;
-              }
-
-              RemoteNode remote = connectionHandler.getRemote();
+            .channelCloseAction(remote -> {
               node.moduleLookup.remove(remote);
               node.moduleConnection.remove(remote);
             })
diff --git a/encon/src/main/java/io/appulse/encon/connection/handshake/AbstractHandshakeChannelInitializer.java b/encon/src/main/java/io/appulse/encon/connection/handshake/AbstractHandshakeChannelInitializer.java
@@ -22,7 +22,6 @@
 import static lombok.AccessLevel.PROTECTED;
 
 import io.netty.channel.ChannelDuplexHandler;
-import io.netty.channel.ChannelFutureListener;
 import io.netty.channel.ChannelInboundHandler;
 import io.netty.channel.ChannelInitializer;
 import io.netty.channel.ChannelOutboundHandler;
@@ -32,7 +31,6 @@
 import io.netty.handler.codec.LengthFieldPrepender;
 import io.netty.handler.logging.LoggingHandler;
 import io.netty.handler.timeout.ReadTimeoutHandler;
-
 import lombok.NonNull;
 import lombok.RequiredArgsConstructor;
 import lombok.experimental.FieldDefaults;
@@ -76,9 +74,6 @@ static void cleanup (@NonNull ChannelPipeline pipeline) {
   @NonNull
   ChannelInboundHandler decoder;
 
-  @NonNull
-  ChannelFutureListener channelCloseListener;
-
   @Override
   protected void initChannel (SocketChannel socketChannel) throws Exception {
     throw new UnsupportedOperationException();
@@ -94,8 +89,6 @@ protected void initChannel (SocketChannel socketChannel, AbstractHandshakeHandle
         .addLast("ENCODER", ENCODER)
         .addLast("HANDLER", handler);
 
-    socketChannel.closeFuture().addListener(channelCloseListener);
-
     log.debug("Handshake pipeline for {} was initialized",
               socketChannel.remoteAddress());
   }
diff --git a/encon/src/main/java/io/appulse/encon/connection/handshake/AbstractHandshakeHandler.java b/encon/src/main/java/io/appulse/encon/connection/handshake/AbstractHandshakeHandler.java
@@ -19,6 +19,7 @@
 import static lombok.AccessLevel.PROTECTED;
 
 import java.util.concurrent.CompletableFuture;
+import java.util.function.Consumer;
 
 import io.appulse.encon.Node;
 import io.appulse.encon.common.RemoteNode;
@@ -51,6 +52,9 @@ abstract class AbstractHandshakeHandler extends ChannelInboundHandlerAdapter {
   @NonNull
   CompletableFuture<Connection> future;
 
+  @NonNull
+  Consumer<RemoteNode> channelCloseAction;
+
   @NonFinal
   RemoteNode remote;
 
@@ -74,7 +78,7 @@ protected void successHandshake (@NonNull Channel channel) {
     log.debug("Replacing pipline to regular for {}", channel.remoteAddress());
 
     AbstractHandshakeChannelInitializer.cleanup(pipeline);
-    val handler = RegularPipeline.setup(pipeline, node, remote);
+    val handler = RegularPipeline.setup(pipeline, node, remote, channelCloseAction);
 
     future.complete(new Connection(remote, handler));
 
diff --git a/encon/src/main/java/io/appulse/encon/connection/handshake/HandshakeClientInitializer.java b/encon/src/main/java/io/appulse/encon/connection/handshake/HandshakeClientInitializer.java
@@ -19,12 +19,12 @@
 import static lombok.AccessLevel.PRIVATE;
 
 import java.util.concurrent.CompletableFuture;
+import java.util.function.Consumer;
 
 import io.appulse.encon.Node;
 import io.appulse.encon.common.RemoteNode;
 import io.appulse.encon.connection.Connection;
 
-import io.netty.channel.ChannelFutureListener;
 import io.netty.channel.ChannelInboundHandler;
 import io.netty.channel.socket.SocketChannel;
 import lombok.Builder;
@@ -52,21 +52,24 @@ public final class HandshakeClientInitializer extends AbstractHandshakeChannelIn
 
   RemoteNode remote;
 
+  Consumer<RemoteNode> channelCloseAction;
+
   @Builder
   public HandshakeClientInitializer (@NonNull Node node,
                                      @NonNull CompletableFuture<Connection> future,
                                      @NonNull RemoteNode remote,
-                                     ChannelFutureListener channelCloseListener
+                                     @NonNull Consumer<RemoteNode> channelCloseAction
   ) {
-    super(DECODER, channelCloseListener);
+    super(DECODER);
     this.node = node;
     this.future = future;
     this.remote = remote;
+    this.channelCloseAction = channelCloseAction;
   }
 
   @Override
   protected void initChannel (SocketChannel socketChannel) throws Exception {
-    val handler = new HandshakeHandlerClient(node, future, remote);
+    val handler = new HandshakeHandlerClient(node, future, remote, channelCloseAction);
     initChannel(socketChannel, handler);
   }
 }
diff --git a/encon/src/main/java/io/appulse/encon/connection/handshake/HandshakeHandlerClient.java b/encon/src/main/java/io/appulse/encon/connection/handshake/HandshakeHandlerClient.java
@@ -21,6 +21,7 @@
 import java.util.Arrays;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ThreadLocalRandom;
+import java.util.function.Consumer;
 
 import io.appulse.encon.Node;
 import io.appulse.encon.common.RemoteNode;
@@ -54,8 +55,12 @@ class HandshakeHandlerClient extends AbstractHandshakeHandler {
   int myChallenge;
 
   @Builder
-  HandshakeHandlerClient (Node node, CompletableFuture<Connection> future, @NonNull RemoteNode remote) {
-    super(node, future);
+  HandshakeHandlerClient (Node node,
+                          CompletableFuture<Connection> future,
+                          @NonNull RemoteNode remote,
+                          Consumer<RemoteNode> channelCloseAction
+  ) {
+    super(node, future, channelCloseAction);
     this.remote = remote;
   }
 
diff --git a/encon/src/main/java/io/appulse/encon/connection/handshake/HandshakeHandlerServer.java b/encon/src/main/java/io/appulse/encon/connection/handshake/HandshakeHandlerServer.java
@@ -22,8 +22,10 @@
 import java.util.Arrays;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ThreadLocalRandom;
+import java.util.function.Consumer;
 
 import io.appulse.encon.Node;
+import io.appulse.encon.common.RemoteNode;
 import io.appulse.encon.connection.Connection;
 import io.appulse.encon.connection.handshake.exception.HandshakeException;
 import io.appulse.encon.connection.handshake.message.ChallengeAcknowledgeMessage;
@@ -54,8 +56,8 @@ class HandshakeHandlerServer extends AbstractHandshakeHandler {
   int ourChallenge;
 
   @Builder
-  HandshakeHandlerServer (Node node, CompletableFuture<Connection> future) {
-    super(node, future);
+  HandshakeHandlerServer (Node node, CompletableFuture<Connection> future, Consumer<RemoteNode> channelCloseAction) {
+    super(node, future, channelCloseAction);
   }
 
   @Override
diff --git a/encon/src/main/java/io/appulse/encon/connection/handshake/HandshakeServerInitializer.java b/encon/src/main/java/io/appulse/encon/connection/handshake/HandshakeServerInitializer.java
@@ -22,9 +22,9 @@
 import java.util.function.Consumer;
 
 import io.appulse.encon.Node;
+import io.appulse.encon.common.RemoteNode;
 import io.appulse.encon.connection.Connection;
 
-import io.netty.channel.ChannelFutureListener;
 import io.netty.channel.ChannelInboundHandler;
 import io.netty.channel.socket.SocketChannel;
 import lombok.Builder;
@@ -52,14 +52,17 @@ public final class HandshakeServerInitializer extends AbstractHandshakeChannelIn
 
   Consumer<CompletableFuture<Connection>> consumer;
 
+  Consumer<RemoteNode> channelCloseAction;
+
   @Builder
   public HandshakeServerInitializer (@NonNull Node node,
                                      @NonNull Consumer<CompletableFuture<Connection>> consumer,
-                                     ChannelFutureListener channelCloseListener
+                                     @NonNull Consumer<RemoteNode> channelCloseAction
   ) {
-    super(DECODER, channelCloseListener);
+    super(DECODER);
     this.node = node;
     this.consumer = consumer;
+    this.channelCloseAction = channelCloseAction;
   }
 
   @Override
@@ -68,7 +71,7 @@ protected void initChannel (SocketChannel socketChannel) throws Exception {
               socketChannel.remoteAddress());
 
     CompletableFuture<Connection> future = new CompletableFuture<>();
-    val handler = new HandshakeHandlerServer(node, future);
+    val handler = new HandshakeHandlerServer(node, future, channelCloseAction);
     initChannel(socketChannel, handler);
     consumer.accept(future);
   }
diff --git a/encon/src/main/java/io/appulse/encon/connection/regular/ConnectionHandler.java b/encon/src/main/java/io/appulse/encon/connection/regular/ConnectionHandler.java
@@ -19,6 +19,7 @@
 import static lombok.AccessLevel.PRIVATE;
 
 import java.io.Closeable;
+import java.util.function.Consumer;
 
 import io.appulse.encon.Node;
 import io.appulse.encon.common.RemoteNode;
@@ -59,6 +60,9 @@ public final class ConnectionHandler extends ChannelInboundHandlerAdapter implem
   @NonNull
   RemoteNode remote;
 
+  @NonNull
+  Consumer<RemoteNode> channelCloseAction;
+
   @NonFinal
   Channel channel;
 
@@ -122,6 +126,8 @@ public void close () {
     if (channel.isOpen()) {
       channel.close();
     }
+    channelCloseAction.accept(remote);
+
     log.debug("Client handler for {} was closed", channel.remoteAddress());
   }
 
diff --git a/encon/src/main/java/io/appulse/encon/connection/regular/RegularPipeline.java b/encon/src/main/java/io/appulse/encon/connection/regular/RegularPipeline.java
@@ -19,6 +19,8 @@
 import static io.netty.handler.logging.LogLevel.DEBUG;
 import static java.lang.Integer.MAX_VALUE;
 
+import java.util.function.Consumer;
+
 import io.appulse.encon.Node;
 import io.appulse.encon.common.RemoteNode;
 
@@ -59,9 +61,10 @@ public final class RegularPipeline {
 
   public static ConnectionHandler setup (@NonNull ChannelPipeline pipeline,
                                          @NonNull Node node,
-                                         @NonNull RemoteNode remoteNode
+                                         @NonNull RemoteNode remoteNode,
+                                         @NonNull Consumer<RemoteNode> channelCloseAction
   ) {
-    val handler = new ConnectionHandler(node, remoteNode);
+    val handler = new ConnectionHandler(node, remoteNode, channelCloseAction);
 
     pipeline
         .addLast(LOGGING_HANDLER)
diff --git a/encon/src/test/java/io/appulse/encon/NodeTest.java b/encon/src/test/java/io/appulse/encon/NodeTest.java
@@ -105,8 +105,8 @@ public void ping () throws Exception {
     val name1 = createFullName();
     val name2 = createFullName();
     node = Nodes.singleNode(name1, NodeConfig.builder()
-                            .cookie("secret")
-                            .build()
+        .cookie("secret")
+        .build()
     );
 
     assertThat(node.ping(ELIXIR_ECHO_SERVER))
