backport FeedableBodyGenerator from master

Author: yanns

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

@@ -0,0 +1,117 @@
+package com.ning.http.client.providers.netty;
+
+import com.ning.http.client.Body;
+import com.ning.http.client.BodyGenerator;
+
+import java.io.IOException;
+import java.io.UnsupportedEncodingException;
+import java.nio.ByteBuffer;
+import java.util.Queue;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * {@link com.ning.http.client.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 static final String US_ASCII = "US-ASCII";
+    private final static byte[] END_PADDING = getBytes("\r\n");
+    private final static byte[] ZERO = getBytes("0");
+    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 {
+        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(getBytes(Integer.toHexString(size)));
+            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;
+        }
+    }
+
+    private static byte[] getBytes(String s) {
+        // for compatibility with java5, we cannot use s.getBytes(Charset)
+        try {
+            return s.getBytes(US_ASCII);
+        } catch (UnsupportedEncodingException e) {
+            throw new RuntimeException(e);
+        }
+    }
+}

src/test/java/com/ning/http/client/async/netty/NettyFeedableBodyGeneratorTest.java

@@ -0,0 +1,128 @@
+/*
+ * Copyright (c) 2013-2014 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.async.netty;
+
+import com.ning.http.client.*;
+import com.ning.http.client.async.AbstractBasicTest;
+import com.ning.http.client.async.ChunkingTest;
+import com.ning.http.client.async.ProviderUtil;
+import com.ning.http.client.providers.netty.FeedableBodyGenerator;
+import org.eclipse.jetty.server.handler.AbstractHandler;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+import javax.servlet.ServletException;
+import javax.servlet.ServletInputStream;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.net.URL;
+import java.nio.ByteBuffer;
+import java.nio.channels.FileChannel;
+
+import static org.testng.FileAssert.fail;
+
+public class NettyFeedableBodyGeneratorTest extends AbstractBasicTest {
+
+    @Override
+    public AsyncHttpClient getAsyncHttpClient(AsyncHttpClientConfig config) {
+        return ProviderUtil.nettyProvider(config);
+    }
+
+    @Test(groups = { "standalone", "default_provider" }, enabled = true)
+    public void testPutImageFile() throws Exception {
+        File largeFile = getTestFile();
+        final FileChannel fileChannel = new FileInputStream(largeFile).getChannel();
+
+        AsyncHttpClientConfig config = new AsyncHttpClientConfig.Builder().setRequestTimeoutInMs(100 * 6000).build();
+        AsyncHttpClient client = getAsyncHttpClient(config);
+        final FeedableBodyGenerator bodyGenerator = new FeedableBodyGenerator();
+
+        try {
+            RequestBuilder builder = new RequestBuilder("PUT")
+                    .setUrl(getTargetUrl())
+                    .setBody(bodyGenerator);
+
+            ListenableFuture<Response> listenableFuture = client.executeRequest(builder.build());
+
+            boolean repeat = true;
+            while (repeat) {
+                final ByteBuffer buffer = ByteBuffer.allocate(1024);
+                if (fileChannel.read(buffer) > 0) {
+                    buffer.flip();
+                    bodyGenerator.feed(buffer, false);
+                } else {
+                    repeat = false;
+                }
+            }
+            ByteBuffer emptyBuffer = ByteBuffer.wrap(new byte[0]);
+            bodyGenerator.feed(emptyBuffer, true);
+
+            Response response = listenableFuture.get();
+            Assert.assertEquals(200, response.getStatusCode());
+            Assert.assertEquals("" + largeFile.length(), response.getHeader("X-TRANSFERRED"));
+        } finally {
+            fileChannel.close();
+            client.close();
+        }
+    }
+
+    private static File getTestFile() {
+        String testResource1 = "300k.png";
+
+        File testResource1File = null;
+        try {
+            ClassLoader cl = ChunkingTest.class.getClassLoader();
+            URL url = cl.getResource(testResource1);
+            testResource1File = new File(url.toURI());
+        } catch (Throwable e) {
+            // TODO Auto-generated catch block
+            fail("unable to find " + testResource1);
+        }
+
+        return testResource1File;
+    }
+
+    @Override
+    public AbstractHandler configureHandler() throws Exception {
+        return new AbstractHandler() {
+
+            public void handle(String arg0, org.eclipse.jetty.server.Request arg1, HttpServletRequest req, HttpServletResponse resp) throws IOException, ServletException {
+
+                ServletInputStream in = req.getInputStream();
+                byte[] b = new byte[8192];
+
+                int count = -1;
+                int total = 0;
+                while ((count = in.read(b)) != -1) {
+                    b = new byte[8192];
+                    total += count;
+                }
+                System.err.println("consumed " + total + " bytes.");
+
+                resp.setStatus(200);
+                resp.addHeader("X-TRANSFERRED", String.valueOf(total));
+                resp.getOutputStream().flush();
+                resp.getOutputStream().close();
+
+                arg1.setHandled(true);
+
+            }
+        };
+    }
+
+
+}