Repository: netty/netty
Issue: Small race in ChunkedWriteHandler
Labels: defect

We have a small race in ChunkedWriteHandler as currentMessageEvent is not "marked" volatile. This needs to get changed


Issue comments:

Pull request review comments:

Patch that resolved the issue:
diff --git a/src/main/java/org/jboss/netty/handler/stream/ChunkedWriteHandler.java b/src/main/java/org/jboss/netty/handler/stream/ChunkedWriteHandler.java
index 96b2838511b..1e714ed3b0d 100644
--- a/src/main/java/org/jboss/netty/handler/stream/ChunkedWriteHandler.java
+++ b/src/main/java/org/jboss/netty/handler/stream/ChunkedWriteHandler.java
@@ -19,6 +19,7 @@
 
 import java.nio.channels.ClosedChannelException;
 import java.util.Queue;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.jboss.netty.buffer.ChannelBuffers;
 import org.jboss.netty.channel.Channel;
@@ -78,6 +79,9 @@ public class ChunkedWriteHandler implements ChannelUpstreamHandler, ChannelDowns
     private final Queue<MessageEvent> queue = QueueFactory.createQueue(MessageEvent.class);
 
     private volatile ChannelHandlerContext ctx;
+    
+    private final AtomicBoolean flush = new AtomicBoolean(false);
+    
     private MessageEvent currentEvent;
 
     /**
@@ -116,12 +120,11 @@ public void handleDownstream(ChannelHandlerContext ctx, ChannelEvent e)
         assert offered;
 
         final Channel channel = ctx.getChannel();
-        if (channel.isWritable()) {
+        
+        // call flush if the channel is writable or not connected. flush(..) will take care of the rest
+        if (channel.isWritable() || !channel.isConnected()) {
             this.ctx = ctx;
             flush(ctx, false);
-        } else if (!channel.isConnected()) {
-            this.ctx = ctx;
-            discard(ctx, false);
         }
     }
 
@@ -147,7 +150,6 @@ public void handleUpstream(ChannelHandlerContext ctx, ChannelEvent e)
 
     private void discard(ChannelHandlerContext ctx, boolean fireNow) {
         ClosedChannelException cause = null;
-        boolean fireExceptionCaught = false;
            
         for (;;) {
             MessageEvent currentEvent = this.currentEvent;
@@ -173,13 +175,12 @@ private void discard(ChannelHandlerContext ctx, boolean fireNow) {
                 cause = new ClosedChannelException();
             }
             currentEvent.getFuture().setFailure(cause);
-            fireExceptionCaught = true;
 
             currentEvent = null;
         }
         
 
-        if (fireExceptionCaught) {
+        if (cause != null) {
             if (fireNow) {
                 Channels.fireExceptionCaught(ctx.getChannel(), cause);
             } else {
@@ -188,95 +189,110 @@ private void discard(ChannelHandlerContext ctx, boolean fireNow) {
         }
     }
 
-    private synchronized void flush(ChannelHandlerContext ctx, boolean fireNow) throws Exception {
-        final Channel channel = ctx.getChannel();
-        if (!channel.isConnected()) {
-            discard(ctx, fireNow);
-        }
+    private void flush(ChannelHandlerContext ctx, boolean fireNow) {
+        boolean acquired = false;
+        
+        // use CAS to see if the have flush already running, if so we don't need to take futher actions
+        if (acquired = flush.compareAndSet(false, true)) {
+            try {
+                
+                final Channel channel = ctx.getChannel();
+                if (!channel.isConnected()) {
+                    discard(ctx, fireNow);
+                }
 
-        while (channel.isWritable()) {
-            if (currentEvent == null) {
-                currentEvent = queue.poll();
-            }
+                while (channel.isWritable()) {
+                    if (currentEvent == null) {
+                        currentEvent = queue.poll();
+                    }
+                    if (currentEvent == null) {
+                        break;
+                    }
 
-            if (currentEvent == null) {
-                break;
-            }
+                    if (currentEvent.getFuture().isDone()) {
+                        // Skip the current request because the previous partial write
+                        // attempt for the current request has been failed.
+                        currentEvent = null;
+                    } else {
+                        final MessageEvent currentEvent = this.currentEvent;
+                        Object m = currentEvent.getMessage();
+                        if (m instanceof ChunkedInput) {
+                            ChunkedInput chunks = (ChunkedInput) m;
+                            Object chunk;
+                            boolean endOfInput;
+                            boolean suspend;
+                            try {
+                                chunk = chunks.nextChunk();
+                                endOfInput = chunks.isEndOfInput();
+                                if (chunk == null) {
+                                    chunk = ChannelBuffers.EMPTY_BUFFER;
+                                    // No need to suspend when reached at the end.
+                                    suspend = !endOfInput;
+                                } else {
+                                    suspend = false;
+                                }
+                            } catch (Throwable t) {
+                                this.currentEvent = null;
 
-            if (currentEvent.getFuture().isDone()) {
-                // Skip the current request because the previous partial write
-                // attempt for the current request has been failed.
-                currentEvent = null;
-            } else {
-                final MessageEvent currentEvent = this.currentEvent;
-                Object m = currentEvent.getMessage();
-                if (m instanceof ChunkedInput) {
-                    ChunkedInput chunks = (ChunkedInput) m;
-                    Object chunk;
-                    boolean endOfInput;
-                    boolean suspend;
-                    try {
-                        chunk = chunks.nextChunk();
-                        endOfInput = chunks.isEndOfInput();
-                        if (chunk == null) {
-                            chunk = ChannelBuffers.EMPTY_BUFFER;
-                            // No need to suspend when reached at the end.
-                            suspend = !endOfInput;
-                        } else {
-                            suspend = false;
-                        }
-                    } catch (Throwable t) {
-                        this.currentEvent = null;
+                                currentEvent.getFuture().setFailure(t);
+                                if (fireNow) {  
+                                    fireExceptionCaught(ctx, t);
+                                } else {
+                                    fireExceptionCaughtLater(ctx, t);
+                                }
 
-                        currentEvent.getFuture().setFailure(t);
-                        if (fireNow) {  
-                            fireExceptionCaught(ctx, t);
-                        } else {
-                            fireExceptionCaughtLater(ctx, t);
-                        }
+                                closeInput(chunks);
+                                break;
+                            }
 
-                        closeInput(chunks);
-                        break;
-                    }
+                            if (suspend) {
+                                // ChunkedInput.nextChunk() returned null and it has
+                                // not reached at the end of input.  Let's wait until
+                                // more chunks arrive.  Nothing to write or notify.
+                                break;
+                            } else {
+                                ChannelFuture writeFuture;
+                                if (endOfInput) {
+                                    this.currentEvent = null;
+                                    closeInput(chunks);
+                                    writeFuture = currentEvent.getFuture();
+                                } else {
+                                    writeFuture = future(channel);
+                                    writeFuture.addListener(new ChannelFutureListener() {
+                                        public void operationComplete(ChannelFuture future)
+                                                throws Exception {
+                                            if (!future.isSuccess()) {
+                                                currentEvent.getFuture().setFailure(future.getCause());
+                                                closeInput((ChunkedInput) currentEvent.getMessage());
+                                            }
+                                        }
+                                    });
+                                }
 
-                    if (suspend) {
-                        // ChunkedInput.nextChunk() returned null and it has
-                        // not reached at the end of input.  Let's wait until
-                        // more chunks arrive.  Nothing to write or notify.
-                        break;
-                    } else {
-                        ChannelFuture writeFuture;
-                        if (endOfInput) {
-                            this.currentEvent = null;
-                            closeInput(chunks);
-                            writeFuture = currentEvent.getFuture();
+                                Channels.write(
+                                        ctx, writeFuture, chunk,
+                                        currentEvent.getRemoteAddress());
+                            }
                         } else {
-                            writeFuture = future(channel);
-                            writeFuture.addListener(new ChannelFutureListener() {
-                                public void operationComplete(ChannelFuture future)
-                                        throws Exception {
-                                    if (!future.isSuccess()) {
-                                        currentEvent.getFuture().setFailure(future.getCause());
-                                        closeInput((ChunkedInput) currentEvent.getMessage());
-                                    }
-                                }
-                            });
+                            this.currentEvent = null;
+                            ctx.sendDownstream(currentEvent);
                         }
+                    }
 
-                        Channels.write(
-                                ctx, writeFuture, chunk,
-                                currentEvent.getRemoteAddress());
+                    if (!channel.isConnected()) {
+                        discard(ctx, fireNow);
+                        break;
                     }
-                } else {
-                    this.currentEvent = null;
-                    ctx.sendDownstream(currentEvent);
                 }
+            } finally {
+                // mark the flush as done
+                flush.set(false);
             }
-
-            if (!channel.isConnected()) {
-                discard(ctx, fireNow);
-                break;
-            }
+        }
+        
+        Channel channel = ctx.getChannel();
+        if (acquired && !channel.isConnected() || (channel.isWritable() && !queue.isEmpty())) {
+            flush(ctx, fireNow);
         }
     }
 


Test patch: