Non blocking upload for Netty

Author: erwan

providers/netty/src/main/java/com/ning/http/client/providers/netty/BodyChunkedInput.java

@@ -33,20 +33,30 @@ class BodyChunkedInput
 
     private static final ByteBuffer EOF = ByteBuffer.allocate(0);
 
+    private boolean endOfInput = false;
+
     public BodyChunkedInput(Body body) {
         if (body == null) {
             throw new IllegalArgumentException("no body specified");
         }
         this.body = body;
     }
 
-    private ByteBuffer peekNextChuck()
+    private ByteBuffer peekNextChunk()
             throws IOException {
 
         if (nextChunk == null) {
             ByteBuffer buffer = ByteBuffer.allocate(chunkSize);
-            if (body.read(buffer) < 0) {
-                nextChunk = EOF;
+            long length = body.read(buffer);
+            if (length < 0) {
+                // Negative means this is finished
+                buffer.flip();
+                nextChunk = buffer;
+                endOfInput = true;
+            } else if (length == 0) {
+                // Zero means we didn't get anything this time, but may get next time
+                buffer.flip();
+                nextChunk = null;
             } else {
                 buffer.flip();
                 nextChunk = buffer;
@@ -55,28 +65,28 @@ private ByteBuffer peekNextChuck()
         return nextChunk;
     }
 
-    public boolean hasNextChunk()
-            throws Exception {
-        return !isEndOfInput();
+    /**
+     * Having no next chunk does not necessarily means end of input, other chunks may arrive later
+     */
+    public boolean hasNextChunk() throws Exception {
+        return peekNextChunk() != null;
     }
 
-    public Object nextChunk()
-            throws Exception {
-        ByteBuffer buffer = peekNextChuck();
-        if (buffer == EOF) {
+    public Object nextChunk() throws Exception {
+        ByteBuffer buffer = peekNextChunk();
+        if (buffer == null || buffer == EOF) {
             return null;
         }
         nextChunk = null;
+
         return ChannelBuffers.wrappedBuffer(buffer);
     }
 
-    public boolean isEndOfInput()
-            throws Exception {
-        return peekNextChuck() == EOF;
+    public boolean isEndOfInput() throws Exception {
+        return endOfInput;
     }
 
-    public void close()
-            throws Exception {
+    public void close() throws Exception {
         body.close();
     }
 

providers/netty/src/main/java/com/ning/http/client/providers/netty/FeedableBodyGenerator.java

@@ -0,0 +1,119 @@
+/*
+ * Copyright (c) 2012 Sonatype, Inc. All rights reserved.
+ *
+ * This program is licensed to you under the Apache License Version 2.0,
+ * and you may not use this file except in compliance with the Apache License Version 2.0.
+ * You may obtain a copy of the Apache License Version 2.0 at http://www.apache.org/licenses/LICENSE-2.0.
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the Apache License Version 2.0 is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the Apache License Version 2.0 for the specific language governing permissions and limitations there under.
+ */
+package com.ning.http.client.providers.netty;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Queue;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import com.ning.http.client.Body;
+import com.ning.http.client.BodyGenerator;
+
+/**
+ * {@link BodyGenerator} which may return just part of the payload at the time
+ * handler is requesting it. If it happens - PartialBodyGenerator becomes responsible
+ * for finishing payload transferring asynchronously.
+ */
+public class FeedableBodyGenerator implements BodyGenerator {
+    private final static byte[] END_PADDING = "\r\n".getBytes();
+    private final static byte[] ZERO = "0".getBytes();
+    private final Queue<BodyPart> queue = new ConcurrentLinkedQueue<BodyPart>();
+    private final AtomicInteger queueSize = new AtomicInteger();
+    private FeedListener listener;
+
+    @Override
+    public Body createBody() throws IOException {
+        return new PushBody();
+    }
+
+    public void feed(final ByteBuffer buffer, final boolean isLast) throws IOException {
+        queue.offer(new BodyPart(buffer, isLast));
+        queueSize.incrementAndGet();
+        if (listener != null) {
+            listener.onContentAdded();
+        }
+    }
+
+    public static interface FeedListener {
+        public void onContentAdded();
+    }
+
+    public void setListener(FeedListener listener) {
+        this.listener = listener;
+    }
+
+    private final class PushBody implements Body {
+        private final int ONGOING = 0;
+        private final int CLOSING = 1;
+        private final int FINISHED = 2;
+
+        private int finishState = 0;
+
+        @Override
+        public long getContentLength() {
+            return -1;
+        }
+
+        @Override
+        public long read(final ByteBuffer buffer) throws IOException {
+            BodyPart nextPart = queue.peek();
+            if (nextPart == null) {
+                // Nothing in the queue
+                switch (finishState) {
+                case ONGOING:
+                    return 0;
+                case CLOSING:
+                    buffer.put(ZERO);
+                    buffer.put(END_PADDING);
+                    finishState = FINISHED;
+                    return buffer.position();
+                case FINISHED:
+                    buffer.put(END_PADDING);
+                    return -1;
+                }
+            }
+            int capacity = buffer.remaining() - 10; // be safe (we'll have to add size, ending, etc.)
+            int size = Math.min(nextPart.buffer.remaining(), capacity);
+            buffer.put(Integer.toHexString(size).getBytes());
+            buffer.put(END_PADDING);
+            for (int i=0; i < size; i++) {
+              buffer.put(nextPart.buffer.get());
+            }
+            buffer.put(END_PADDING);
+            if (!nextPart.buffer.hasRemaining()) {
+                if (nextPart.isLast) {
+                    finishState = CLOSING;
+                }
+                queue.remove();
+            }
+            return size;
+        }
+
+        @Override
+        public void close() throws IOException {
+        }
+
+    }
+
+    private final static class BodyPart {
+        private final boolean isLast;
+        private final ByteBuffer buffer;
+
+        public BodyPart(final ByteBuffer buffer, final boolean isLast) {
+            this.buffer = buffer;
+            this.isLast = isLast;
+        }
+    }
+}

providers/netty/src/main/java/com/ning/http/client/providers/netty/NettyAsyncHttpProvider.java

@@ -45,6 +45,7 @@
 import com.ning.http.client.listener.TransferCompletionHandler;
 import com.ning.http.client.ntlm.NTLMEngine;
 import com.ning.http.client.ntlm.NTLMEngineException;
+import com.ning.http.client.providers.netty.FeedableBodyGenerator.FeedListener;
 import com.ning.http.client.providers.netty.spnego.SpnegoEngine;
 import com.ning.http.client.websocket.WebSocketUpgradeHandler;
 import com.ning.http.multipart.MultipartBody;
@@ -488,6 +489,14 @@ protected final <T> void writeRequest(final Channel channel,
                             writeFuture = channel.write(bodyFileRegion);
                         } else {
                             BodyChunkedInput bodyChunkedInput = new BodyChunkedInput(body);
+                            BodyGenerator bg = future.getRequest().getBodyGenerator();
+                                if (bg instanceof FeedableBodyGenerator) {
+                                    ((FeedableBodyGenerator)bg).setListener(new FeedListener() {
+                                        @Override public void onContentAdded() {
+                                            channel.getPipeline().get(ChunkedWriteHandler.class).resumeTransfer();
+                                        }
+                                    });
+                                }
                             writeFuture = channel.write(bodyChunkedInput);
                         }