Merge pull request #54 from rogierslag/reproducer/delayed-jackson-response-lifecycle

Wait for pending body writes before completing Netty responses

Author: jamezp

embedded-server/pom.xml

@@ -115,6 +115,11 @@
             <artifactId>resteasy-jaxb-provider</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.jboss.resteasy</groupId>
+            <artifactId>resteasy-jackson2-provider</artifactId>
+            <scope>test</scope>
+        </dependency>
         <!-- This needs a new release of RESTEasy 6.2.13+ for this to be included
         <dependency>
             <groupId>org.jboss.resteasy</groupId>

embedded-server/src/main/java/org/jboss/resteasy/plugins/server/netty/ChunkOutputStream.java

@@ -6,6 +6,7 @@
 package org.jboss.resteasy.plugins.server.netty;
 
 import java.io.IOException;
+import java.io.OutputStream;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 
@@ -14,9 +15,11 @@
 
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
+import io.netty.channel.ChannelFuture;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelPromise;
 import io.netty.handler.codec.http.DefaultHttpContent;
+import io.netty.util.concurrent.Future;
 
 /**
  * Class to help application that are built to write to an
@@ -43,6 +46,12 @@ public class ChunkOutputStream extends AsyncOutputStream {
     private final ByteBuf buffer;
     private final ChannelHandlerContext ctx;
     private final NettyHttpResponse response;
+    // All lifecycle state below is guarded by writeLock.
+    private int pendingWrites = 0;
+    private Throwable writeFailure = null;
+    private ChannelPromise responsePromise = null;
+    private boolean finishRequested = false;
+    private boolean responseWriteStarted = false;
 
     ChunkOutputStream(final NettyHttpResponse response, final ChannelHandlerContext ctx, final int chunksize) {
         this.response = response;
@@ -56,6 +65,7 @@ public class ChunkOutputStream extends AsyncOutputStream {
     @Override
     public void write(int b) throws IOException {
         synchronized (writeLock) {
+            ensureOpen();
             if (buffer.maxWritableBytes() < 1) {
                 flush();
             }
@@ -67,6 +77,8 @@ public void reset() {
         if (response.isCommitted())
             throw new IllegalStateException(Messages.MESSAGES.responseIsCommitted());
         synchronized (writeLock) {
+            if (finishRequested)
+                throw new IllegalStateException(Messages.MESSAGES.responseIsCommitted());
             buffer.clear();
         }
     }
@@ -88,6 +100,7 @@ private void write(byte[] b, int off, int len, ChannelPromise promise) throws IO
         int spaceLeftInCurrentChunk;
         MultiPromise mp = new MultiPromise(ctx, promise);
         synchronized (writeLock) {
+            ensureOpen();
             while ((spaceLeftInCurrentChunk = buffer.maxWritableBytes()) < dataLengthLeftToWrite) {
                 buffer.writeBytes(b, dataToWriteOffset, spaceLeftInCurrentChunk);
                 dataToWriteOffset = dataToWriteOffset + spaceLeftInCurrentChunk;
@@ -109,13 +122,16 @@ public void flush() throws IOException {
 
     private void flush(ChannelPromise promise) throws IOException {
         synchronized (writeLock) {
+            ensureOpen();
             int readable = buffer.readableBytes();
             if (readable == 0) {
                 promise.setSuccess();
                 return;
             }
             if (!response.isCommitted())
                 response.prepareChunkStream();
+            pendingWrites++;
+            promise.addListener(this::bodyWriteComplete);
             ctx.writeAndFlush(new DefaultHttpContent(buffer.copy()), promise);
             buffer.clear();
         }
@@ -157,4 +173,85 @@ public CompletionStage<Void> asyncWrite(byte[] bytes, int offset, int length) {
         }
         return ret;
     }
+
+    /**
+     * Closes the entity-output lifecycle and completes the HTTP response after every body write has completed.
+     * The entity stream can wrap this root stream, hence it is flushed while holding the write lock. Any tail emitted by
+     * that flush is registered before the response is marked as finished.
+     */
+    ChannelFuture finish(OutputStream entityOutputStream) throws IOException {
+        ChannelPromise result;
+        boolean completeResponse;
+        synchronized (writeLock) {
+            if (finishRequested) {
+                return responsePromise;
+            }
+            if (entityOutputStream != null) {
+                entityOutputStream.flush();
+            }
+            finishRequested = true;
+            responsePromise = ctx.newPromise();
+            result = responsePromise;
+            completeResponse = pendingWrites == 0;
+        }
+        if (completeResponse) {
+            completeResponse();
+        }
+        return result;
+    }
+
+    private void ensureOpen() throws IOException {
+        if (finishRequested) {
+            throw new IOException(Messages.MESSAGES.responseIsCommitted());
+        }
+    }
+
+    private void bodyWriteComplete(Future<?> future) {
+        boolean completeResponse;
+        synchronized (writeLock) {
+            pendingWrites--;
+            if (!future.isSuccess() && writeFailure == null) {
+                writeFailure = future.cause() == null
+                        ? new IOException("Response body write failed without a cause")
+                        : future.cause();
+            }
+            completeResponse = finishRequested && pendingWrites == 0;
+        }
+        if (!future.isSuccess()) {
+            // Once body bytes may have reached the peer, only closing the transport can prevent a clean partial response.
+            ctx.close();
+        }
+        if (completeResponse) {
+            completeResponse();
+        }
+    }
+
+    private void completeResponse() {
+        final ChannelPromise result;
+        final Throwable failure;
+        synchronized (writeLock) {
+            if (!finishRequested || pendingWrites != 0 || responseWriteStarted) {
+                return;
+            }
+            responseWriteStarted = true;
+            result = responsePromise;
+            failure = writeFailure;
+        }
+
+        if (failure != null) {
+            ctx.close().addListener(ignored -> result.tryFailure(failure));
+            return;
+        }
+
+        response.writeResponseTermination().addListener(future -> {
+            if (future.isSuccess()) {
+                result.trySuccess();
+            } else {
+                Throwable cause = future.cause() == null
+                        ? new IOException("Response termination write failed without a cause")
+                        : future.cause();
+                ctx.close().addListener(ignored -> result.tryFailure(cause));
+            }
+        });
+    }
 }

embedded-server/src/main/java/org/jboss/resteasy/plugins/server/netty/NettyHttpResponse.java

@@ -32,19 +32,28 @@
 import io.netty.handler.codec.http.LastHttpContent;
 
 /**
+ * Body write submission and completion are ordered by the root {@link ChunkOutputStream}'s write lock. Terminal writes
+ * can be requested by RESTEasy dispatch or by an asynchronous Netty promise callback. {@link #sendError(int, String)}
+ * and {@link #writeResponseTermination()} therefore synchronize on this response, making terminal message selection,
+ * write submission and future publication one transition.
+ *
  * @author <a href="mailto:bill@burkecentral.com">Bill Burke</a>
  * @version $Revision: 1 $
  */
 public class NettyHttpResponse implements HttpResponse {
     private static final int EMPTY_CONTENT_LENGTH = 0;
     private int status = 200;
-    private OutputStream os;
+    // RESTEasy writer interceptors can replace this stream with a wrapper around the root stream.
+    private OutputStream entityOutputStream;
+    // The root stream owns the Netty body-write promises and therefore remains stable when the entity stream is replaced.
+    private final ChunkOutputStream rootChunkOutputStream;
     private final MultivaluedMap<String, Object> outputHeaders;
     private final ChannelHandlerContext ctx;
     private boolean committed;
     private final boolean keepAlive;
     private final ResteasyProviderFactory providerFactory;
     private final HttpMethod method;
+    private ChannelFuture terminationFuture;
 
     public NettyHttpResponse(final ChannelHandlerContext ctx, final boolean keepAlive,
             final ResteasyProviderFactory providerFactory) {
@@ -55,15 +64,18 @@ public NettyHttpResponse(final ChannelHandlerContext ctx, final boolean keepAliv
             final ResteasyProviderFactory providerFactory, final HttpMethod method) {
         outputHeaders = new MultivaluedMapImpl<String, Object>();
         this.method = method;
-        os = (method == null || !method.equals(HttpMethod.HEAD)) ? new ChunkOutputStream(this, ctx, 1000) : null; //[RESTEASY-1627]
+        rootChunkOutputStream = (method == null || !method.equals(HttpMethod.HEAD))
+                ? new ChunkOutputStream(this, ctx, 1000)
+                : null; //[RESTEASY-1627]
+        entityOutputStream = rootChunkOutputStream;
         this.ctx = ctx;
         this.keepAlive = keepAlive;
         this.providerFactory = providerFactory;
     }
 
     @Override
-    public void setOutputStream(OutputStream os) {
-        this.os = os;
+    public void setOutputStream(OutputStream entityOutputStream) {
+        this.entityOutputStream = entityOutputStream;
     }
 
     @Override
@@ -83,7 +95,7 @@ public MultivaluedMap<String, Object> getOutputHeaders() {
 
     @Override
     public OutputStream getOutputStream() throws IOException {
-        return os;
+        return entityOutputStream;
     }
 
     @Override
@@ -97,7 +109,7 @@ public void sendError(int status) throws IOException {
     }
 
     @Override
-    public void sendError(int status, String message) throws IOException {
+    public synchronized void sendError(int status, String message) throws IOException {
         if (committed) {
             throw new IllegalStateException();
         }
@@ -116,8 +128,8 @@ public void sendError(int status, String message) throws IOException {
         }
         // Add keep alive or connection close header
         transformResponseHeaders(response);
-        ctx.writeAndFlush(response);
         committed = true;
+        terminationFuture = ctx.writeAndFlush(response);
     }
 
     @Override
@@ -157,6 +169,10 @@ private void transformResponseHeaders(io.netty.handler.codec.http.HttpResponse r
         RestEasyHttpResponseEncoder.transformHeaders(this, res, providerFactory);
     }
 
+    /**
+     * Called by {@link ChunkOutputStream} while it holds its write lock. Body-write completion reacquires that lock before
+     * it can request termination, which makes this committed state visible to the completion path.
+     */
     public void prepareChunkStream() {
         committed = true;
         DefaultHttpResponse response = getDefaultHttpResponse();
@@ -165,14 +181,11 @@ public void prepareChunkStream() {
     }
 
     public void finish() throws IOException {
-        if (os != null)
-            os.flush();
         ChannelFuture future;
-        if (isCommitted()) {
-            // if committed this means the output stream was used.
-            future = ctx.writeAndFlush(LastHttpContent.EMPTY_LAST_CONTENT);
+        if (rootChunkOutputStream != null) {
+            future = rootChunkOutputStream.finish(entityOutputStream);
         } else {
-            future = ctx.writeAndFlush(getEmptyHttpResponse());
+            future = writeResponseTermination();
         }
 
         if (!isKeepAlive()) {
@@ -181,10 +194,20 @@ public void finish() throws IOException {
 
     }
 
+    synchronized ChannelFuture writeResponseTermination() {
+        if (terminationFuture != null) {
+            return terminationFuture;
+        }
+        Object terminalMessage = isCommitted() ? LastHttpContent.EMPTY_LAST_CONTENT : getEmptyHttpResponse();
+        committed = true;
+        terminationFuture = ctx.writeAndFlush(terminalMessage);
+        return terminationFuture;
+    }
+
     @Override
     public void flushBuffer() throws IOException {
-        if (os != null)
-            os.flush();
+        if (entityOutputStream != null)
+            entityOutputStream.flush();
         ctx.flush();
     }
 }