Issue #16: Data can be sent to a wrong subscriber when pipeline is enabled

alaisi · alaisi · commit f0c7eae4bcbb · 2015-11-26T21:50:12.000+02:00
diff --git a/src/main/java/com/github/pgasync/impl/netty/NettyPgConnectionPool.java b/src/main/java/com/github/pgasync/impl/netty/NettyPgConnectionPool.java
@@ -39,15 +39,17 @@ public class NettyPgConnectionPool extends PgConnectionPool {
 
     final EventLoopGroup group = new NioEventLoopGroup(1);
     final boolean useSsl;
+    final boolean pipeline;
 
     public NettyPgConnectionPool(PoolProperties properties) {
         super(properties);
         useSsl = properties.getUseSsl();
+        pipeline = properties.getUsePipelining();
     }
 
     @Override
     protected PgProtocolStream openStream(InetSocketAddress address) {
-        return new NettyPgProtocolStream(group, address, useSsl);
+        return new NettyPgProtocolStream(group, address, useSsl, pipeline);
     }
 
     @Override
diff --git a/src/main/java/com/github/pgasync/impl/netty/NettyPgProtocolStream.java b/src/main/java/com/github/pgasync/impl/netty/NettyPgProtocolStream.java
@@ -46,19 +46,23 @@
 public class NettyPgProtocolStream implements PgProtocolStream {
 
     final EventLoopGroup group;
+    final EventLoop eventLoop;
     final SocketAddress address;
     final boolean useSsl;
+    final boolean pipeline;
 
     final GenericFutureListener<Future<? super Object>> onError;
     final Queue<Subscriber<? super Message>> subscribers;
     final ConcurrentMap<String,Map<String,Subscriber<? super String>>> listeners = new ConcurrentHashMap<>();
 
     ChannelHandlerContext ctx;
 
-    public NettyPgProtocolStream(EventLoopGroup group, SocketAddress address, boolean useSsl) {
+    public NettyPgProtocolStream(EventLoopGroup group, SocketAddress address, boolean useSsl, boolean pipeline) {
         this.group = group;
+        this.eventLoop = group.next();
         this.address = address;
         this.useSsl = useSsl; // TODO: refactor into SSLConfig with trust parameters
+        this.pipeline = pipeline;
         this.subscribers = new LinkedBlockingDeque<>(); // TODO: limit pipeline queue depth
         this.onError = future -> {
             if(!future.isSuccess()) {
@@ -101,6 +105,14 @@ public Observable<Message> send(Message... messages) {
                 return;
             }
 
+            if(pipeline && !eventLoop.inEventLoop()) {
+                eventLoop.submit(() -> {
+                    pushSubscriber(subscriber);
+                    write(messages);
+                });
+                return;
+            }
+
             pushSubscriber(subscriber);
             write(messages);
 
@@ -136,14 +148,16 @@ public Observable<String> listen(String channel) {
 
     @Override
     public Observable<Void> close() {
-        return Observable.create(subscriber -> ctx.close().addListener(f -> {
-            if(!f.isSuccess()) {
-                subscriber.onError(f.cause());
+        return Observable.create(subscriber ->
+                    ctx.writeAndFlush(Terminate.INSTANCE).addListener(written ->
+                            ctx.close().addListener(closed -> {
+            if (!closed.isSuccess()) {
+                subscriber.onError(closed.cause());
                 return;
             }
             subscriber.onNext(null);
             subscriber.onCompleted();
-        }));
+        })));
     }
 
     private void pushSubscriber(Subscriber<? super Message> subscriber) {
