Catch errors from writing to netty context

alaisi · alaisi · commit c54ce97bf675 · 2015-09-19T13:18:28.000+03:00
diff --git a/src/main/java/com/github/pgasync/impl/netty/NettyPgProtocolStream.java b/src/main/java/com/github/pgasync/impl/netty/NettyPgProtocolStream.java
@@ -26,17 +26,12 @@
 import io.netty.handler.ssl.SslContext;
 import io.netty.handler.ssl.SslHandshakeCompletionEvent;
 import io.netty.handler.ssl.util.InsecureTrustManagerFactory;
+import io.netty.util.concurrent.*;
+import io.netty.util.concurrent.Future;
 
 import java.net.SocketAddress;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Map;
-import java.util.UUID;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
-import java.util.Queue;
-import java.util.concurrent.ArrayBlockingQueue;
-import java.util.concurrent.LinkedBlockingDeque;
+import java.util.*;
+import java.util.concurrent.*;
 import java.util.function.Consumer;
 
 import static java.util.Collections.singletonList;
@@ -70,11 +65,7 @@ public void connect(StartupMessage startup, Consumer<List<Message>> replyTo) {
                 .channel(NioSocketChannel.class)
                 .handler(newProtocolInitializer(newStartupHandler(startup, replyTo)))
                 .connect(address)
-                .addListener(future -> {
-                    if (!future.isSuccess()) {
-                        replyTo.accept(singletonList(new ChannelError(future.cause())));
-                    }
-                });
+                .addListener(errorListener(replyTo));
     }
 
     @Override
@@ -83,7 +74,7 @@ public void send(Message message, Consumer<List<Message>> replyTo) {
             throw new IllegalStateException("Channel is closed");
         }
         addNewReplyHandler(replyTo);
-        ctx.writeAndFlush(message);
+        ctx.writeAndFlush(message).addListener(errorListener(replyTo));
     }
 
     private void addNewReplyHandler(Consumer<List<Message>> replyTo) {
@@ -98,7 +89,8 @@ public void send(List<Message> messages, Consumer<List<Message>> replyTo) {
             throw new IllegalStateException("Channel is closed");
         }
         addNewReplyHandler(replyTo);
-        messages.forEach(ctx::write);
+        GenericFutureListener<Future<Object>> errorListener = errorListener(replyTo);
+        messages.forEach(msg -> ctx.write(msg).addListener(errorListener));
         ctx.flush();
     }
 
@@ -151,6 +143,14 @@ void publishNotification(NotificationResponse notification) {
         }
     }
 
+    <T> GenericFutureListener<io.netty.util.concurrent.Future<T>> errorListener(Consumer<List<Message>> replyTo) {
+        return future -> {
+            if(!future.isSuccess()) {
+                replyTo.accept(singletonList(new ChannelError(future.cause())));
+            }
+        };
+    }
+
     ChannelInboundHandlerAdapter newStartupHandler(StartupMessage startup, Consumer<List<Message>> replyTo) {
         return new ChannelInboundHandlerAdapter() {
             @Override
@@ -162,15 +162,15 @@ public void userEventTriggered(ChannelHandlerContext context, Object evt) throws
             @Override
             public void channelActive(ChannelHandlerContext context) {
                 if(useSsl) {
-                    context.writeAndFlush(SSLHandshake.INSTANCE);
+                    context.writeAndFlush(SSLHandshake.INSTANCE).addListener(errorListener(replyTo));
                 } else {
                     startup(context);
                 }
             }
             void startup(ChannelHandlerContext context) {
                 ctx = context;
                 addNewReplyHandler(replyTo);
-                context.writeAndFlush(startup);
+                context.writeAndFlush(startup).addListener(errorListener(replyTo));
                 context.pipeline().remove(this);
             }
         };
